import os, pathlib, asyncio from time import time try: from telegram import ForceReply, Update from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters, ConversationHandler except ModuleNotFoundError: Update = None class ct: DEFAULT_TYPE = None ContextTypes = ct() print("Cannot import telegram futures") filter_tables = { "v1":{#default "resize_filter":"[logo][0]scale2ref=w=oh*mdar:h=ih*1.0[logo][0];", "insert_filter":"[0][logo]overlay=(W-w)/2:(H-h)/2:format=auto,format=yuv420p", "filter":"[1]format=rgba,colorchannelmixer=aa={transparency}[logo];{resize_filter}{insert_filter}" }, "v2":{#photo "filter":"[1]format=rgba,colorchannelmixer=aa={transparency}[logo];[logo][0]scale2ref=oh*mdar:ih*0.2[logo][video];[video][logo]overlay=(main_w-overlay_w):(main_h-overlay_h)" }, "v2.move":{#good "insert_filter":"[video][logo]overlay=x='if(lt(mod(t\,16)\,8)\,W-w-W*10/100\,W*10/100)':y='if(lt(mod(t+4\,16)\,8)\,H-h-H*5/100\,H*5/100)'", "filter":"[1]colorchannelmixer=aa={transparency}[logo];[logo][0]scale2ref=oh*mdar:ih*0.2[logo][video];{insert_filter}" }, "v2.dvd":{ "insert_filter":"[video][logo]overlay=x='if(eq(mod(floor(st(8,t*{speed_x})/(W-w)),2),0),mod(ld(8),W-w),W-w-mod(ld(8),W-w))':y='if(eq(mod(floor(st(9,t*{speed_y})/(H-h)),2),0),mod(ld(9),H-h),H-h-mod(ld(9),H-h))'", "filter":"[1]colorchannelmixer=aa={transparency}[logo];[logo][0]scale2ref=oh*mdar:ih*0.2[logo][video];{insert_filter}" } } async def add_watermark(input_file, watermark, transparency, out_ext, filter_version = "v1", speed=(160, 70)): resize_filter = filter_tables[filter_version].get("resize_filter", "") insert_filter = filter_tables[filter_version].get("insert_filter", "").format(speed_x = speed[0], speed_y = speed[1]) filter = filter_tables[filter_version]["filter"].format(transparency = transparency, resize_filter = resize_filter, insert_filter = insert_filter) cmd = f'ffmpeg -y -i {input_file} -i {watermark} -filter_complex "{filter}" {"-c:v libx264 -c:a copy" if out_ext == "mp4" else ""} {input_file}.{out_ext}' print(cmd) process = await asyncio.create_subprocess_shell(cmd) await process.wait() if os.path.exists(f"{input_file}.{out_ext}"): return f"{input_file}.{out_ext}" return "" class Extension: cmds = ["/watermark"] wait_content = {} content_directory = os.getenv("CONTENT_DIR") if os.getenv("CONTENT_DIR", None) else "./content" png_path = os.getenv("WATERMARK") default_transparency=20.0 watermark_users = [] def __init__(self, core) -> None: self.watermark_users = core.admins print(f"Content save path: {self.content_directory}") print(f"Watermark users: {' '.join([str(i) for i in self.watermark_users])}") async def __call__(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if update.message.text and update.message.text.split()[0] in self.cmds: return await self.watermark(update, context) return await self.content_handler(update, context) async def watermark(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if not update.message.from_user.id in self.watermark_users: return await update.message.reply_text("Это не для таких как ты сделано, друг...") if update.message.from_user.id in self.wait_content: del self.wait_content[update.message.from_user.id] return await update.message.reply_text("Понял понял, не бей, отменяем думанье...") else: try: self.wait_content[update.message.from_user.id] = float(update.message.text.split()[1]) except: self.wait_content[update.message.from_user.id] = self.default_transparency return await update.message.reply_text(f"Отлично, жду от тебя ОДНУ фотку/видео/переслал\nНепрозрачность: {self.wait_content[update.message.from_user.id]}%") async def content_handler(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if not update.message.from_user.id in self.wait_content: return if update.message.video: return await self.video_handler(update, context) elif update.message.photo: return await self.photo_handler(update, context) #Что ниже над будет засунуть в одно, ахуй async def photo_handler(self, update:Update, context: ContextTypes.DEFAULT_TYPE): if update.message.from_user.id in self.wait_content: transparency = self.wait_content[update.message.from_user.id] / 100 del self.wait_content[update.message.from_user.id] else: return _id = update.message.photo[-1].file_id photo = await context.bot.get_file(_id) downloaded_photo = await photo.download_to_drive(f"{self.content_directory}/{time()}") res = await add_watermark(downloaded_photo.absolute(), self.png_path, transparency, "jpg", "v2") if res: await update.message.reply_photo(photo=res) os.remove(res) else: await update.message.reply_text("Ошибка при создании фотки") if os.path.exists(downloaded_photo.absolute()): os.remove(downloaded_photo.absolute()) return async def video_handler(self, update: Update, context: ContextTypes.DEFAULT_TYPE): if update.message.from_user.id in self.wait_content: transparency = self.wait_content[update.message.from_user.id] / 100 del self.wait_content[update.message.from_user.id] else: return video = await update.message.video.get_file() downloaded_video = await video.download_to_drive(f"{self.content_directory}/{time()}") res = await add_watermark(downloaded_video.absolute(), self.png_path, transparency, "mp4", "v2.dvd") if res: await update.message.reply_video(video=res) os.remove(res) else: await update.message.reply_text("Ошибка при создании видео") if os.path.exists(downloaded_video.absolute()): os.remove(downloaded_video.absolute()) return if __name__ == "__main__": import asyncio import argparse parser = argparse.ArgumentParser() parser.add_argument("input", type=str) parser.add_argument("watermark", type=str) parser.add_argument("--transparency", type=float, default=100.0) parser.add_argument("--insert-version", type=str, default="v1") parser.add_argument("--out-ext", type=str, default="mp4") args = parser.parse_args() async def run(): await add_watermark(args.input, args.watermark, args.transparency / 100, args.out_ext, args.insert_version) asyncio.get_event_loop().run_until_complete(run())