warp-beacon 1.0.2__py3-none-any.whl

etc/warp_beacon/warp_beacon.conf

@@ -0,0 +1,10 @@
+TG_TOKEN=""
+INSTAGRAM_LOGIN=""
+INSTAGRAM_PASSWORD=""
+INSTAGRAM_VERIFICATION_CODE=""
+MONGODB_HOST="mongodb"
+MONGODB_PORT="27017"
+MONGODB_USER="root"
+MONGODB_PASSWORD="changeme"
+VIDEO_STORAGE_DIR="/var/warp_beacon/videos"
+WORKERS_POOL_SIZE=3

lib/systemd/system/warp_beacon.service

@@ -0,0 +1,14 @@
+[Unit]
+Description=Telegram bot for expanding external media links
+After=network-online.target syslog.target network.target remote-fs.target nss-lookup.target multi-user.target
+
+[Service]
+User=root
+Group=root
+Type=simple
+Restart=always
+EnvironmentFile=/etc/warp_beacon/warp_beacon.conf
+ExecStart=/opt/venvs/warp-beacon/bin/warp_beacon
+
+[Install]
+WantedBy=multi-user.target

warp_beacon/__version__.py

@@ -0,0 +1,2 @@
+__version__ = "1.0.2"
+

warp_beacon/jobs/abstract.py

@@ -0,0 +1,60 @@
+from abc import ABC, abstractmethod
+from typing import TypedDict
+from typing_extensions import Unpack
+import uuid
+
+class JobSettings(TypedDict):
+	job_id: uuid.UUID
+	message_id: int
+	local_media_path: str
+	media_info: dict
+	url: str
+	uniq_id: str
+	tg_file_id: str
+	in_process: bool
+	job_failed: bool
+	media_type: str
+	job_failed_msg: str
+	effective_url: str
+	save_items: bool
+	media_collection: list
+
+class AbstractJob(ABC):
+	job_id: uuid.UUID = None
+	message_id: int = 0
+	local_media_path: str = ""
+	media_info: dict = {}
+	url: str = ""
+	uniq_id: str = ""
+	tg_file_id: str = ""
+	media_type: str = "video"
+	in_process: bool = False
+	job_failed: bool = False
+	job_failed_msg: str = ""
+	effective_url: str = ""
+	save_items: bool = False
+	media_collection: list = []
+
+	def __init__(self, **kwargs: Unpack[JobSettings]) -> None:
+		if kwargs:
+			self.__dict__.update(kwargs)
+		self.job_id = uuid.uuid4()
+
+	def __del__(self) -> None:
+		pass
+
+	def __str__(self) -> str:
+		return str(self.to_dict())
+
+	def __repr__(self) -> str:
+		return str(self.to_dict())
+
+	def to_dict(self) -> dict:
+		d = {}
+		for key in dir(self.__class__):
+			if not key.startswith('_'):
+				value = getattr(self, key)
+				if not callable(value):
+					d[key] = value
+					
+		return d

warp_beacon/jobs/download_job.py

@@ -0,0 +1,23 @@
+from typing import TypedDict
+from typing_extensions import Unpack
+
+from warp_beacon.jobs.upload_job import UploadJob
+from warp_beacon.jobs.abstract import AbstractJob, JobSettings
+
+class DownloadJob(AbstractJob):
+	def __init__(self, **kwargs: Unpack[JobSettings]) -> None:
+		super(DownloadJob, self).__init__(**kwargs)
+	def __del__(self) -> None:
+		pass
+
+	@staticmethod
+	def build(**kwargs: Unpack[JobSettings]) -> "DownloadJob":
+		return DownloadJob(**kwargs)
+	
+	def to_upload_job(self, **kwargs: Unpack[JobSettings]) -> AbstractJob:
+		d = self.to_dict()
+		d.update(kwargs)
+		if "media_collection" in d:
+			for k, v in enumerate(d["media_collection"]):
+				d["media_collection"][k] = UploadJob.build(**v)
+		return UploadJob.build(**d)

warp_beacon/jobs/upload_job.py

@@ -0,0 +1,26 @@
+from typing import TypedDict
+from typing_extensions import Unpack
+
+from warp_beacon.jobs.abstract import AbstractJob, JobSettings
+
+class UploadJob(AbstractJob):
+	def __init__(self, **kwargs: Unpack[JobSettings]) -> None:
+		super(UploadJob, self).__init__(**kwargs)
+	def __del__(self) -> None:
+		pass
+
+	@staticmethod
+	def build(**kwargs: Unpack[JobSettings]) -> "UploadJob":
+		return UploadJob(**kwargs)
+	
+	def to_download_job(self, **kwargs: Unpack[JobSettings]) -> AbstractJob:
+		from warp_beacon.jobs.download_job import DownloadJob
+		d = self.to_dict()
+		d.update(kwargs)
+		return DownloadJob.build(**d)
+	
+	def set_flag(self, key: str, value: bool) -> "UploadJob":
+		if key in self.__dict__:
+			self.__dict__[key] = value
+
+		return self

warp_beacon/mediainfo/video.py

@@ -0,0 +1,80 @@
+import io, os
+from typing import Optional
+import cv2
+
+class VideoInfo(object):
+	vid = None
+	# need for filesize
+	filename = ""
+
+	def __init__(self, filename: str) -> None:
+		self.vid = cv2.VideoCapture(filename)
+		self.filename = filename
+		
+	def __del__(self) -> None:
+		self.vid.release()
+
+	def get_demensions(self) -> dict:
+		res = {"width": None, "height": None}
+		if self.vid.isOpened():
+			res["width"] = int(self.vid.get(cv2.CAP_PROP_FRAME_WIDTH))
+			res["height"] = int(self.vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+		return res
+
+	def get_duration(self) -> int:
+		duration_in_seconds = None
+		if self.vid.isOpened():
+			fps = self.vid.get(cv2.CAP_PROP_FPS)
+			total_no_frames = self.vid.get(cv2.CAP_PROP_FRAME_COUNT)
+			duration_in_seconds = int(total_no_frames / fps)
+
+		return duration_in_seconds
+	
+	def get_filesize(self) -> float:
+		size = os.path.getsize(self.filename)
+		return round(size/(pow(1024,2)), 2)
+	
+	def get_finfo(self, except_info: tuple=()) -> dict:
+		res = {}
+		res.update(self.get_demensions())
+		if "duration" not in except_info:
+			res["duration"] = self.get_duration()
+		if "filesize" not in except_info:
+			res["filesize"] = self.get_filesize()
+		return res
+	
+	def shrink_image_to_fit(self, img):
+		height, width = img.shape[:2]
+		max_height = 320
+		max_width = 320
+
+		# only shrink if img is bigger than required
+		if max_height < height or max_width < width:
+			# get scaling factor
+			scaling_factor = max_height / float(height)
+			if max_width/float(width) < scaling_factor:
+				scaling_factor = max_width / float(width)
+			# resize image
+			img = cv2.resize(img, None, fx=scaling_factor, fy=scaling_factor, interpolation=cv2.INTER_AREA)
+
+		return img
+	
+	def generate_thumbnail(self) -> Optional[io.BytesIO]:
+		if self.vid.isOpened():
+			count = 4
+			success = True
+			while success:
+				self.vid.set(cv2.CAP_PROP_POS_MSEC,(count*1000))
+				success, image = self.vid.read()
+				if success:
+					image = self.shrink_image_to_fit(image)
+					success, buffer = cv2.imencode(".jpg", image)
+				if success:
+					io_buf = io.BytesIO(buffer)
+					io_buf.seek(0)
+					#io_buf.name = "thumbnail.png"
+					return io_buf
+				count += 1
+
+		return None

warp_beacon/scrapler/__init__.py

@@ -0,0 +1,155 @@
+from typing import Optional
+import multiprocessing
+import time
+import logging
+from requests.exceptions import ConnectTimeout, HTTPError
+from instagrapi.exceptions import MediaNotFound, UnknownError
+
+from warp_beacon.mediainfo.video import VideoInfo
+from warp_beacon.uploader import AsyncUploader
+from warp_beacon.jobs.download_job import DownloadJob
+
+CONST_CPU_COUNT = multiprocessing.cpu_count()
+
+class AsyncDownloader(object):
+	workers = []
+	allow_loop = None
+	job_queue = multiprocessing.Queue()
+	uploader = None
+	workers_count = CONST_CPU_COUNT
+
+	def __init__(self, uploader: AsyncUploader, workers_count: int=CONST_CPU_COUNT) -> None:
+		self.allow_loop = multiprocessing.Value('i', 1)
+		self.uploader = uploader
+		self.workers_count = workers_count
+
+	def __del__(self) -> None:
+		self.stop_all()
+
+	def start(self) -> None:
+		for _ in range(self.workers_count):
+			proc = multiprocessing.Process(target=self.do_work)
+			self.workers.append(proc)
+			proc.start()
+
+	def get_media_info(self, path: str, fr_media_info: dict={}) -> Optional[dict]:
+		media_info = None
+		try:
+			if path:
+				video_info = VideoInfo(path)
+				media_info = video_info.get_finfo(tuple(fr_media_info.keys()))
+				media_info.update(fr_media_info)
+				logging.info("Media file info: %s", media_info)
+				media_info["thumb"] = video_info.generate_thumbnail()
+		except Exception as e:
+			logging.error("Failed to process media info!")
+			logging.exception(e)
+
+		return media_info
+
+	def do_work(self) -> None:
+		logging.info("download worker started")
+		while self.allow_loop.value == 1:
+			try:
+				job = None
+				try:
+					job = self.job_queue.get()
+					actor = None
+					try:
+						items = []
+						if "instagram.com/" in job.url:
+							if not job.in_process:
+								from warp_beacon.scrapler.instagram import InstagramScrapler
+								actor = InstagramScrapler()
+								while True:
+									try:
+										logging.info("Downloading URL '%s'", job.url)
+										items = actor.download(job.url)
+										break
+									except ConnectTimeout as e:
+										logging.error("ConnectTimeout download error!")
+										logging.exception(e)
+										time.sleep(2)
+									except MediaNotFound as e:
+										logging.warning("MediaNotFound occurred!")
+										logging.exception(e)
+										self.uploader.queue_task(job.to_upload_job(
+											job_failed=True,
+											job_failed_msg="Unable to access to media under this URL. Seems like the media is private.")
+										)
+										break
+									except (UnknownError, Exception) as e:
+										logging.warning("UnknownError occurred!")
+										logging.exception(e)
+										exception_msg = ""
+										if hasattr(e, "message"):
+											exception_msg = e.message
+										else:
+											exception_msg = str(e)
+										if "geoblock_required" in exception_msg:
+											self.uploader.queue_task(job.to_upload_job(
+												job_failed=True,
+												job_failed_msg="This content does not accessible for bot account. Seems like author blocked certain region.")
+											)
+											break
+										self.uploader.queue_task(job.to_upload_job(
+											job_failed=True,
+											job_failed_msg="WOW, unknown error occured! Please send service logs to developer via email: andrey@bagrintsev.me.")
+										)
+										break
+
+								if items:
+									for item in items:
+										media_info = {"filesize": 0}
+										if item["media_type"] == "video":
+											media_info = self.get_media_info(item["local_media_path"], item["media_info"])
+										elif item["media_type"] == "collection":
+											for v in item["items"]:
+												if v["media_type"] == "video":
+													col_media_info = self.get_media_info(v["local_media_path"], v["media_info"])
+													media_info["filesize"] += int(col_media_info.get("filesize", 0))
+													v["media_info"] = col_media_info
+
+										job_args = {"media_type": item["media_type"], "media_info": media_info}
+										if item["media_type"] == "collection":
+											job_args["media_collection"] = item["items"]
+											if item.get("save_items", None) is not None:
+												job_args["save_items"] = item.get("save_items", False)
+										else:
+											job_args["local_media_path"] = item["local_media_path"]
+
+										upload_job = job.to_upload_job(**job_args)
+										self.uploader.queue_task(upload_job)
+							else:
+								logging.info("Job already in work in parallel worker. Redirecting job to upload worker.")
+								self.uploader.queue_task(job.to_upload_job())
+					except HTTPError as e:
+						logging.error("HTTP error inside download worker!")
+						logging.exception(e)
+					except Exception as e:
+						logging.error("Error inside download worker!")
+						logging.exception(e)
+						self.notify_task_failed(job)
+						#self.queue_task(url=item["url"], message_id=item["message_id"], item_in_process=item["in_process"], uniq_id=item["uniq_id"])
+				except multiprocessing.Queue.empty:
+					pass
+			except Exception as e:
+				logging.error("Exception occurred inside worker!")
+				logging.exception(e)
+
+	def stop_all(self) -> None:
+		self.allow_loop.value = 0
+		for proc in self.workers:
+			if proc.is_alive():
+				logging.info("stopping process #%d", proc.pid)
+				proc.terminate()
+				proc.join()
+				logging.info("process #%d stopped", proc.pid)
+		self.workers.clear()
+
+	def queue_task(self, job: DownloadJob) -> str:
+		self.job_queue.put_nowait(job)
+		return str(job.job_id)
+	
+	def notify_task_failed(self, job: DownloadJob) -> None:
+		self.uploader.queue_task(job.to_upload_job(job_failed=True))

warp_beacon/scrapler/abstract.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+
+class ScraplerAbstract(ABC):
+	def __init__(self) -> None:
+		pass
+
+	def __del__(self) -> None:
+		pass
+
+	@abstractmethod
+	def scrap(self, url: str) -> str:
+		raise NotImplementedError
+
+	@abstractmethod
+	def download(self, url: str) -> bool:
+		raise NotImplementedError

warp_beacon/scrapler/instagram.py

@@ -0,0 +1,191 @@
+import os
+from pathlib import Path
+import time
+import json
+from typing import Callable, Optional, Union
+
+import requests
+import urllib3
+import logging
+
+from instagrapi.mixins.story import Story
+from instagrapi.types import Media
+from instagrapi import Client
+from instagrapi.exceptions import LoginRequired, PleaseWaitFewMinutes
+
+from warp_beacon.scrapler.abstract import ScraplerAbstract
+
+INST_SESSION_FILE = "/var/warp_beacon/inst_session.json"
+
+class InstagramScrapler(ScraplerAbstract):
+	cl = None
+
+	def __init__(self) -> None:
+		super().__init__()
+		self.cl = Client()
+
+	def safe_write_session(self) -> None:
+		tmp_fname = "%s~" % INST_SESSION_FILE
+		with open(tmp_fname, 'w+') as f:
+			f.write(json.dumps(self.cl.get_settings()))
+		if os.path.isfile(INST_SESSION_FILE):
+			os.unlink(INST_SESSION_FILE)
+		os.rename(tmp_fname, INST_SESSION_FILE)
+
+	def load_session(self) -> None:
+		if os.path.exists(INST_SESSION_FILE):
+			self.cl.load_settings(INST_SESSION_FILE)
+		else:
+			self.login()
+
+	def login(self) -> None:
+		self.cl = Client()
+		username = os.environ.get("INSTAGRAM_LOGIN", default=None)
+		password = os.environ.get("INSTAGRAM_PASSWORD", default=None)
+		verification_code = os.environ.get("INSTAGRAM_VERIFICATION_CODE", default="")
+		if username is not None and password is not None:
+			self.cl.login(username=username, password=password, verification_code=verification_code)
+		self.safe_write_session()
+
+	def scrap(self, url: str) -> tuple[str]:
+		self.load_session()
+		def _scrap() -> tuple[str]:
+			if "stories" in url:
+				url_last_part = list(filter(None, url.split('/')))[-1]
+				if url_last_part.isnumeric():
+					return "story", self.scrap_story(url)
+				else:
+					return "stories", url_last_part
+			else:
+				return "media", self.scrap_media(url)
+		try:
+			return _scrap()
+		except LoginRequired as e:
+			logging.warning("Session error. Trying to relogin...")
+			logging.exception(e)
+			self.login()
+			return _scrap()
+
+	def scrap_stories(self, username: str) -> list[Story]:
+		user_info = self.cl.user_info_by_username(username)
+		logging.info("user_id is '%s'", user_info.pk)
+		return self.cl.user_stories(user_id=user_info.pk)
+
+	def scrap_story(self, url: str) -> str:
+		story_id = self.cl.story_pk_from_url(url)
+		logging.info("story_id is '%s'", story_id)
+		return story_id
+
+	def scrap_media(self, url: str) -> str:
+		media_id = self.cl.media_pk_from_url(url)
+		logging.info("media_id is '%s'", media_id)
+		return media_id
+	
+	def __download_hndlr(self, func: Callable, *args: tuple[str], **kwargs: dict[str]) -> Union[Path, Media]:
+		ret_val = {}
+		max_retries = int(os.environ.get("IG_MAX_RETRIES", default=5))
+		retries = 0
+		while max_retries >= retries:
+			try:
+				ret_val = func(*args, **kwargs)
+				break
+			except (requests.exceptions.ConnectionError,
+					requests.exceptions.ReadTimeout,
+					urllib3.exceptions.ReadTimeoutError,
+					urllib3.exceptions.ConnectionError) as e:
+				logging.warning("Instagram read timeout! Retrying in 2 seconds ...")
+				logging.info("Your `IG_MAX_RETRIES` values is %d", max_retries)
+				logging.exception(e)
+				if max_retries == retries:
+					raise e
+				retries += 1
+				time.sleep(2)
+
+		return ret_val
+
+	
+	def download_video(self, url: str, media_info: dict) -> dict:
+		path = self.__download_hndlr(self.cl.video_download_by_url, url, folder='/tmp')
+		return {"local_media_path": str(path), "media_type": "video", "media_info": {"duration": media_info.video_duration}}
+
+	def download_photo(self, url: str) -> dict:
+		path = self.__download_hndlr(self.cl.photo_download_by_url, url, folder='/tmp')
+		return {"local_media_path": str(path), "media_type": "image"}
+	
+	def download_story(self, story_info: Story) -> dict:
+		path, media_type, media_info = "", "", {}
+		logging.info("Story id is '%s'", story_info.id)
+		effective_story_id = story_info.id
+		if '_' in effective_story_id:
+			st_parts = effective_story_id.split('_')
+			if len(st_parts) > 1:
+				effective_story_id = st_parts[0]
+		logging.info("Effective story id is '%s'", effective_story_id)
+		effective_url = "https://www.instagram.com/stories/%s/%s/" % (story_info.user.username, effective_story_id)
+		if story_info.media_type == 1: # photo
+			path = self.__download_hndlr(self.cl.story_download_by_url, url=story_info.thumbnail_url, folder='/tmp')
+			media_type = "image"
+		elif story_info.media_type == 2: # video
+			path = self.__download_hndlr(self.cl.story_download_by_url, url=story_info.video_url, folder='/tmp')
+			media_type = "video"
+			media_info["duration"] = story_info.video_duration
+
+		return {"local_media_path": str(path), "media_type": media_type, "media_info": media_info, "effective_url": effective_url}
+
+	def download_stories(self, stories: list[Story]) -> dict:
+		res = []
+		for story in stories:
+			res.append(self.download_story(story_info=story))
+
+		return {"media_type": "collection", "save_items": True, "items": res}
+
+	def download_album(self, media_info: dict) -> dict:
+		res = []
+		for i in media_info.resources:
+			_media_info = self.cl.media_info(i.pk)
+			if i.media_type == 1: # photo
+				res.append(self.download_photo(url=_media_info.thumbnail_url))
+			elif i.media_type == 2: # video
+				res.append(self.download_video(url=_media_info.video_url, media_info=_media_info))
+
+		return {"media_type": "collection", "items": res}
+
+	def download(self, url: str) -> Optional[list[dict]]:
+		res = []
+		while True:
+			try:
+				scrap_type, media_id = self.scrap(url)
+				if scrap_type == "media":
+					media_info = self.__download_hndlr(self.cl.media_info, media_id)
+					logging.info("media_type is '%d', product_type is '%s'", media_info.media_type, media_info.product_type)
+					if media_info.media_type == 2 and media_info.product_type == "clips": # Reels
+						res.append(self.download_video(url=media_info.video_url, media_info=media_info))
+					elif media_info.media_type == 1: # Photo
+						res.append(self.download_photo(url=media_info.thumbnail_url))
+					elif media_info.media_type == 8: # Album
+						res.append(self.download_album(media_info=media_info))
+				elif scrap_type == "story":
+					story_info = self.cl.story_info(media_id)
+					logging.info("media_type for story is '%d'", story_info.media_type)
+					res.append(self.download_story(story_info=story_info))
+				elif scrap_type == "stories":
+					logging.info("Stories download mode")
+					res.append(self.download_stories(self.scrap_stories(media_id)))
+				break
+			except PleaseWaitFewMinutes as e:
+				logging.warning("Please wait a few minutes error. Trying to relogin ...")
+				logging.exception(e)
+				wait_timeout = int(os.environ.get("IG_WAIT_TIMEOUT", default=5))
+				logging.info("Waiting %d seconds according configuration option `IG_WAIT_TIMEOUT`", wait_timeout)
+				if res:
+					for i in res:
+						if i["media_type"] == "collection":
+							for j in i["items"]:
+								if os.path.exists(j["local_media_path"]):
+									os.unlink(j["local_media_path"])
+						else:
+							if os.path.exists(i["local_media_path"]):
+								os.unlink(i["local_media_path"])
+				os.unlink(INST_SESSION_FILE)
+				time.sleep(wait_timeout)
+		return res

warp_beacon/storage/__init__.py

@@ -0,0 +1,82 @@
+import os
+#from typing import Optional
+import logging
+
+from urllib.parse import urlparse
+
+from pymongo import MongoClient
+
+VIDEO_STORAGE_DIR = os.environ.get("VIDEO_STORAGE_DIR", default="/var/warp_beacon/videos")
+
+class Storage(object):
+	client = None
+	db = None
+	def __init__(self) -> None:
+		if not os.path.isdir(VIDEO_STORAGE_DIR):
+			os.mkdir(VIDEO_STORAGE_DIR)
+
+		self.client = MongoClient(
+			host=os.environ.get("MONGODB_HOST", default='127.0.0.1'), 
+			port=int(os.environ.get("MONGODB_PORT", default=27017)),
+			username=os.environ.get("MONGODB_USER", default='root'),
+			password=os.environ.get("MONGODB_PASSWORD", default="changeme"))
+		self.db = self.client.media.media
+
+	def __del__(self) -> None:
+		if self.client:
+			self.client.close()
+
+	@staticmethod
+	def compute_uniq(url: str) -> str:
+		path = urlparse(url).path.strip('/')
+		return path
+	
+	def db_find(self, uniq_id: str) -> list[dict]:
+		document = None
+		ret = []
+		try:
+			logging.debug("uniq_id to search is '%s'", uniq_id)
+			cursor = self.db.find({"uniq_id": uniq_id})
+			for document in cursor:
+				ret.append({"uniq_id": document["uniq_id"], "tg_file_id": document["tg_file_id"], "media_type": document["media_type"]})
+		except Exception as e:
+			logging.error("Error occurred while trying to read from the database!")
+			logging.exception(e)
+		return ret
+	
+	def db_lookup(self, url: str) -> dict:
+		uniq_id = self.compute_uniq(url)
+		doc = self.db_find(uniq_id)
+		return doc
+	
+	def db_lookup_id(self, uniq_id: str) -> list[dict]:
+		return self.db_find(uniq_id)
+	
+	def add_media(self, tg_file_ids: list[str], media_url: str, media_type: str, origin: str) -> list[int]:
+		uniq_id = self.compute_uniq(media_url)
+		media_ids = []
+		for tg_file_id in tg_file_ids:
+			if self.db_lookup_id(uniq_id):
+				logging.info("Detected existing uniq_id, skipping storage write operation")
+				continue
+			media_ids += str(self.db.insert_one({"uniq_id": uniq_id, "media_type": media_type, "tg_file_id": tg_file_id, "origin": origin}).inserted_id)
+
+		return media_ids
+	
+	def get_random(self) -> dict:
+		ret = {}
+		try:
+			cursor = self.db.aggregate([
+				{ "$match": { "tg_file_id": { "$exists": True } } },
+				{ "$sample": { "size": 1 } }
+			])
+			tmp = list(cursor)
+			if tmp:
+				ret = tmp.pop()
+		except Exception as e:
+			logging.error("Error occurred while trying to read from the database!")
+			logging.exception(e)
+		return ret
+
+
+