warp-beacon 1.2.6__py3-none-any.whl → 2.0.0__py3-none-any.whl

warp_beacon/telegram/handlers.py

@@ -0,0 +1,156 @@
+from pyrogram import Client
+from pyrogram.types import Message
+from pyrogram.enums import ChatType, ParseMode
+from pyrogram.types import Chat, BotCommand
+
+from urlextract import URLExtract
+
+from warp_beacon.storage import Storage
+from warp_beacon.telegram.utils import Utils
+from warp_beacon.jobs.download_job import DownloadJob
+from warp_beacon.jobs.upload_job import UploadJob
+from warp_beacon.jobs import Origin
+from warp_beacon.jobs.types import JobType
+
+import logging
+
+class Handlers(object):
+	storage = None
+	bot = None
+
+	def __init__(self, bot: "Bot") -> None:
+		self.bot = bot
+		self.storage = bot.storage
+
+	async def help(self, client: Client, message: Message) -> None:
+		"""Send a message when the command /help is issued."""
+		await self.bot.send_text(text="Send me a link to remote media", reply_id=message.id, chat_id=message.chat.id)
+
+	async def random(self, client: Client, message: Message) -> None:
+		d = self.storage.get_random()
+		if not d:
+			await message.reply_text("No random content yet. Try to send link first.")
+			return
+		await self.bot.upload_job(UploadJob(tg_file_id=d["tg_file_id"], media_type=d["media_type"], message_id=message.id))
+
+	async def start(self, client: Client, message: Message) -> None:
+		bot_name = await self.bot.client.get_me()
+		await self.bot.client.set_bot_commands([
+			BotCommand("start", "Start bot"),
+			BotCommand("help", "Show help message"),
+			BotCommand("random", "Get random media")
+		])
+		await message.reply_text(
+			parse_mode=ParseMode.MARKDOWN,
+			text=f"Welcome to @{bot_name.username}!\n"
+			"Send link to external social network with content and I'll reply to it.\n"
+			"Currently supported: Instagram, YouTube Shorts and YouTube Music."
+		)
+
+	async def handler(self, client: Client, message: Message) -> None:
+		chat = message.chat
+		effective_message_id = message.id
+		extractor = URLExtract()
+		urls = extractor.find_urls(message.text)
+
+		reply_text = "Wut?"
+		if not urls:
+			reply_text = "Your message should contains URLs"
+		else:
+			for url in urls:
+				origin = Utils.extract_origin(url)
+				if origin is Origin.YOUTU_BE:
+					url = Utils.extract_youtu_be_link(url)
+					if not url:
+						raise ValueError("Failed to extract youtu.be link")
+					origin = Origin.YOUTUBE
+				if origin is Origin.UNKNOWN:
+					logging.info("Only Instagram, YouTube Shorts and YouTube Music are now supported. Skipping.")
+					continue
+				entities, tg_file_ids = [], []
+				uniq_id = Storage.compute_uniq(url)
+				try:
+					entities = self.storage.db_lookup_id(uniq_id)
+				except Exception as e:
+					logging.error("Failed to search link in DB!")
+					logging.exception(e)
+				if entities:
+					tg_file_ids = [i["tg_file_id"] for i in entities]
+					logging.info("URL '%s' is found in DB. Sending with tg_file_ids = '%s'", url, str(tg_file_ids))
+					ent_len = len(entities)
+					if ent_len > 1:
+						await self.bot.upload_job(
+							UploadJob(
+								tg_file_id=",".join(tg_file_ids),
+								message_id=effective_message_id,
+								media_type=JobType.COLLECTION,
+								chat_id=chat.id
+							)
+						)
+					elif ent_len:
+						logging.info(entities[0])
+						media_type = JobType[entities[0]["media_type"].upper()]
+						await self.bot.upload_job(
+							UploadJob(
+								tg_file_id=tg_file_ids.pop(),
+								message_id=effective_message_id,
+								media_type=media_type,
+								chat_id=chat.id
+							)
+						)
+				else:
+					async def upload_wrapper(job: UploadJob) -> None:
+						try:
+							if job.job_failed and job.job_failed_msg:
+								if job.placeholder_message_id:
+									await self.bot.placeholder.remove(chat.id, job.placeholder_message_id)
+								return await self.bot.send_text(chat_id=chat.id, text=job.job_failed_msg, reply_id=job.message_id)
+							if job.job_warning and job.job_warning_msg:
+								return await self.bot.placeholder.update_text(chat.id, job.placeholder_message_id, job.job_warning_msg)
+							tg_file_ids = await self.bot.upload_job(job)
+							if tg_file_ids:
+								if job.media_type == JobType.COLLECTION and job.save_items:
+									for chunk in job.media_collection:
+										for i in chunk:
+											self.storage.add_media(tg_file_ids=[i.tg_file_id], media_url=i.effective_url, media_type=i.media_type.value, origin=origin.value)
+								else:
+									self.storage.add_media(tg_file_ids=[','.join(tg_file_ids)], media_url=job.url, media_type=job.media_type.value, origin=origin.value)
+						except Exception as e:
+							logging.error("Exception occurred while performing upload callback!")
+							logging.exception(e)
+
+					try:
+						# create placeholder message for long download
+						placeholder_message_id = await self.bot.placeholder.create(
+							chat_id=chat.id,
+							reply_id=effective_message_id
+						)
+
+						if not placeholder_message_id:
+							await self.bot.send_text(
+								chat_id=chat.id,
+								reply_id=effective_message_id, 
+								text="Failed to create message placeholder. Please check your bot Internet connection.")
+							return
+
+						self.bot.uploader.add_callback(
+							placeholder_message_id,
+							upload_wrapper
+						)
+
+						self.bot.downloader.queue_task(DownloadJob.build(
+							url=url,
+							placeholder_message_id=placeholder_message_id,
+							message_id=effective_message_id,
+							chat_id=chat.id,
+							in_process=self.bot.uploader.is_inprocess(uniq_id),
+							uniq_id=uniq_id,
+							job_origin=origin
+						))
+						self.bot.uploader.set_inprocess(uniq_id)
+					except Exception as e:
+						logging.error("Failed to schedule download task!")
+						logging.exception(e)
+
+		if chat.type not in (ChatType.GROUP, ChatType.SUPERGROUP) and not urls:
+			await self.bot.send_text(rext=reply_text, reply_id=effective_message_id)

warp_beacon/telegram/placeholder_message.py

@@ -0,0 +1,191 @@
+import os, io
+import time
+from enum import Enum
+from typing import Optional
+
+from pyrogram.types import Message
+from pyrogram.errors import RPCError, FloodWait
+from pyrogram.enums import ParseMode
+
+import warp_beacon
+from warp_beacon.telegram.utils import Utils
+from warp_beacon.mediainfo.video import VideoInfo
+
+import logging
+
+class PlaceholderType(Enum):
+	UNKNOWN = 0
+	ANIMATION = 1
+	DOCUMENT = 2
+	PHOTO = 3
+
+class PlaceHolder(object):
+	pl_type: PlaceholderType = PlaceholderType.ANIMATION
+	tg_file_id: str = None
+
+	def __init__(self, pl_type: PlaceholderType, tg_file_id: str) -> None:
+		self.pl_type = pl_type
+		self.tg_file_id = tg_file_id
+
+class PlaceholderMessage(object):
+	bot = None
+	placeholder = PlaceHolder(PlaceholderType.ANIMATION, None)
+
+	def __init__(self, bot: "Bot") -> None:
+		self.bot = bot
+
+	def __del__(self) -> None:
+		pass
+
+	async def reuse_ph_animation(self, chat_id: int, reply_id: int, text: str) -> Optional["types.Message"]:
+		reply = await self.bot.client.send_animation(
+			chat_id=chat_id,
+			animation=self.placeholder.tg_file_id,
+			caption=text,
+			parse_mode=ParseMode.MARKDOWN,
+			reply_to_message_id=reply_id
+		)
+
+		return reply
+
+	async def reuse_ph_photo(self, chat_id: int, reply_id: int, text: str) -> Optional["types.Message"]:
+		reply = await self.bot.client.send_photo(
+			chat_id=chat_id,
+			photo=self.placeholder.tg_file_id,
+			caption=text,
+			parse_mode=ParseMode.MARKDOWN,
+			reply_to_message_id=reply_id
+		)
+
+		return reply
+
+	async def reuse_ph_document(self, chat_id: int, reply_id: int, text: str) -> Optional["types.Message"]:
+		reply = await self.bot.client.send_document(
+			chat_id=chat_id,
+			document=self.placeholder.tg_file_id,
+			caption=text,
+			parse_mode=ParseMode.MARKDOWN,
+			reply_to_message_id=reply_id
+		)
+
+		return reply
+
+	async def create(self, chat_id: int, reply_id: int) -> int:
+		retry_amount = 0
+		max_retries = int(os.environ.get("TG_MAX_RETRIES", default=5))
+		while not retry_amount >= max_retries:
+			try:
+				text = "**Loading, this may take a moment ...** ⏱️ "
+				reply = None
+				if self.placeholder.tg_file_id is None:
+					ph_found = False
+					for ph in ('/var/warp_beacon/placeholder.gif', "%s/../var/warp_beacon/placeholder.gif" % os.path.dirname(os.path.abspath(warp_beacon.__file__))):
+						if not os.path.exists(ph):
+							continue
+						try:
+							#pl_info = VideoInfo(ph)
+							#pl_resolution = pl_info.get_demensions()
+							reply = await self.bot.client.send_document(
+								chat_id=chat_id,
+								document=ph,
+								force_document=False,
+								caption=text,
+								parse_mode=ParseMode.MARKDOWN,
+								reply_to_message_id=reply_id,
+								file_name=os.path.basename(ph),
+								#width=pl_resolution["width"],
+								#height=pl_resolution["height"],
+								#duration=round(pl_info.get_duration()),
+								#thumb=pl_info.generate_thumbnail()
+							)
+							self.placeholder = PlaceHolder(PlaceholderType.ANIMATION, Utils.extract_file_id(reply))
+							ph_found = True
+							break
+						except Exception as e:
+							logging.warning("Failed to send placeholder message!")
+							logging.exception(e)
+					if not ph_found:
+						try:
+							reply = await self.bot.client.send_animation(
+								chat_id=chat_id,
+								animation="https://bagrintsev.me/warp_beacon/placeholder_that_we_deserve.mp4",
+								caption=text,
+								parse_mode=ParseMode.MARKDOWN,
+								reply_to_message_id=reply_id
+							)
+							self.placeholder = PlaceHolder(PlaceholderType.ANIMATION, Utils.extract_file_id(reply))
+						except Exception as e:
+							logging.error("Failed to download secret placeholder!")
+							logging.exception(e)
+							img = self.create_default_placeholder_img("Loading, this may take a moment ...")
+							reply = await self.bot.client.send_photo(
+								chat_id=chat.id,
+								parse_mode=ParseMode.MARKDOWN,
+								reply_to_message_id=reply_id,
+								photo=img
+							)
+							self.placeholder = PlaceHolder(PlaceholderType.PHOTO, Utils.extract_file_id(reply))
+				else:
+					if self.placeholder.pl_type == PlaceholderType.ANIMATION:
+						try:
+							reply = await self.reuse_ph_animation(chat_id, reply_id, text)
+						except ValueError as e:
+							logging.warning("Failed to reuse tg_file_id!")
+							logging.exception(e)
+							reply = await self.reuse_ph_document(chat_id, reply_id, text)
+							self.placeholder.pl_type = PlaceholderType.DOCUMENT
+					elif self.placeholder.pl_type == PlaceholderType.DOCUMENT:
+						try:
+							reply = await self.reuse_ph_document(chat_id, reply_id, text)
+						except ValueError as e:
+							logging.warning("Failed to reuse tg_file_id!")
+							logging.exception(e)
+							reply = await self.reuse_ph_animation(chat_id, reply_id, text)
+							self.placeholder.pl_type = PlaceholderType.ANIMATION
+					else:
+						reply = await self.reuse_ph_photo(chat_id, reply_id, text)
+				return reply.id
+			except FloodWait as e:
+				logging.warning("FloodWait exception!")
+				logging.exception(e)
+				await self.bot.send_text(None, "Telegram error: %s" % e.MESSAGE)
+				time.sleep(e.value)
+			except Exception as e:
+				logging.error("Failed to create placeholder message!")
+				logging.exception(e)
+				retry_amount += 1
+				time.sleep(2)
+
+		return 0
+
+	def create_default_placeholder_img(self, text: str, width: int = 800, height: int = 1280) -> io.BytesIO:
+		from PIL import Image, ImageDraw, ImageFont
+		bio = io.BytesIO()
+		bio.name = 'placeholder.png'
+		img = Image.new("RGB", (width, height), (255, 255, 255))
+		draw = ImageDraw.Draw(img)
+		font = ImageFont.load_default(size=48)
+		_, _, w, h = draw.textbbox((0, 0), text, font=font)
+		draw.text(((width-w)/2, (height-h)/2), text, font=font, fill="#000")
+		img.save(bio, 'PNG')
+		bio.seek(0)
+
+		return bio
+
+	async def update_text(self, chat_id: int, placeholder_message_id: int, placeholder_text: str) -> None:
+		try:
+			await self.bot.client.edit_message_caption(
+				chat_id=chat_id,
+				message_id=placeholder_message_id,
+				caption=" ⚠️ *%s*" % placeholder_text
+			)
+		except Exception as e:
+			logging.error("Failed to update placeholder message!")
+			logging.exception(e)
+
+	async def remove(self, chat_id: int, placeholder_message_id: int) -> None:
+		try:
+			await self.bot.client.delete_messages(chat_id, (placeholder_message_id,))
+		except Exception as e:
+			logging.error("Failed to remove placeholder message!")
+			logging.exception(e)

warp_beacon/telegram/utils.py

@@ -0,0 +1,73 @@
+import re
+
+import requests
+
+from enum import Enum
+from typing import Union
+
+from warp_beacon.jobs import Origin
+from warp_beacon.jobs.types import JobType
+
+import logging
+
+class Utils(object):
+	expected_patronum_compiled_re = re.compile(r'Expected ([A-Z]+), got ([A-Z]+) file id instead')
+
+	@staticmethod
+	def extract_file_id(message: "Message") -> Union[None, str]:
+		possible_attrs = ("video", "photo", "audio", "animation", "document")
+		for attr in possible_attrs:
+			if hasattr(message, attr):
+				_attr = getattr(message, attr, None)
+				if _attr:
+					tg_id = getattr(_attr, "file_id", None)
+					if tg_id:
+						return tg_id
+		return None
+
+	@staticmethod
+	def extract_origin(url: str) -> Origin:
+		if "instagram.com/" in url:
+			return Origin.INSTAGRAM
+
+		if "youtube.com/" in url and "shorts/" in url:
+			return Origin.YT_SHORTS
+
+		if "youtube.com/" in url and "music." in url:
+			return Origin.YT_MUSIC
+
+		if "youtu.be/" in url:
+			return Origin.YOUTU_BE
+
+		if "youtube.com/" in url:
+			return Origin.YOUTUBE
+
+		return Origin.UNKNOWN
+
+	@staticmethod
+	def extract_youtu_be_link(url: str) -> str:
+		try:
+			response = requests.get(
+				url=url,
+				allow_redirects=False
+			)
+			return response.headers["Location"]
+		except Exception as e:
+			logging.error("Failed to extract YouTube link!")
+			logging.exception(e)
+
+		return ''
+
+	@staticmethod
+	def parse_expected_patronum_error(err_text: str) -> tuple:
+		'''
+			Input example: 'Expected VIDEO, got ANIMATION file id instead'
+		'''
+		capture = re.match(Utils.expected_patronum_compiled_re, err_text)
+		expected_value, got_value = capture[1], capture[2]
+
+		return JobType[expected_value], JobType[got_value]
+
+	@staticmethod
+	def chunker(seq: list, size: int) -> list:
+		return (seq[pos:pos + size] for pos in range(0, len(seq), size))

warp_beacon/uploader/__init__.py

@@ -5,12 +5,11 @@ from warp_beacon.jobs.upload_job import UploadJob
 import logging
 
 import asyncio
-from telegram import Update
-from telegram.ext import ContextTypes
 
 from typing import Optional, Callable, Coroutine
 
 from warp_beacon.storage import Storage
+from warp_beacon.jobs.types import JobType
 
 class AsyncUploader(object):
 	__JOE_BIDEN_WAKEUP = None
@@ -38,12 +37,12 @@ class AsyncUploader(object):
 			thread.start()
 			self.threads.append(thread)
 
-	def add_callback(self, message_id: int, callback: Callable, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+	def add_callback(self, message_id: int, callback: Callable) -> None:
 		def callback_wrap(*args, **kwargs) -> None:
 			ret = callback(*args, **kwargs)
 			#self.remove_callback(message_id)
 			return ret
-		self.callbacks[message_id] = {"callback": callback_wrap, "update": update, "context": context}
+		self.callbacks[message_id] = {"callback": callback_wrap}
 
 	def remove_callback(self, message_id: int) -> None:
 		if message_id in self.callbacks:
@@ -82,15 +81,16 @@ class AsyncUploader(object):
 					if job is self.__JOE_BIDEN_WAKEUP:
 						continue
 					path = ""
-					if job.media_type == "collection":
+					if job.media_type == JobType.COLLECTION:
 						for i in job.media_collection:
-							path += "%s; " % i.local_media_path
+							for j in i:
+								path += "%s; " % j.local_media_path
 					else:
 						path = job.local_media_path
 					in_process = job.in_process
 					uniq_id = job.uniq_id
 					message_id = job.placeholder_message_id
-					if not in_process and  not job.job_failed and not job.job_warning:
+					if not in_process and not job.job_failed and not job.job_warning:
 						logging.info("Accepted upload job, file(s): '%s'", path)
 					try:
 						if message_id in self.callbacks:
@@ -113,10 +113,10 @@ class AsyncUploader(object):
 									dlds_len = len(db_list_dicts)
 									if dlds_len > 1:
 										job.tg_file_id = ",".join(tg_file_ids)
-										job.media_type = "collection"
+										job.media_type = JobType.COLLECTION
 									elif dlds_len:
 										job.tg_file_id = ",".join(tg_file_ids)
-										job.media_type = db_list_dicts.pop()["media_type"]
+										job.media_type = JobType[db_list_dicts.pop()["media_type"].upper()]
 									asyncio.ensure_future(self.callbacks[message_id]["callback"](job), loop=self.loop)
 									self.process_done(uniq_id)
 									self.remove_callback(message_id)