changes :-

- rewrote tools and its abstractions
- added splitpdf and extactor
- other various changes.

plugins/logger.py → logger.py

@@ -2,13 +2,13 @@
 
 
 logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s - [ %(message)s ]",
- datefmt="%d-%b-%y %H:%M:%S",
- force=True,
+ level=logging.DEBUG,
+ format="%(asctime)s [%(process)d] [%(levelname)s] %(message)s",
  handlers=[logging.FileHandler("debug.log"), logging.StreamHandler()],
+ datefmt="[%Y-%m-%d %H:%M:%S %z]",
+ force=True,
 )
 
 logging.getLogger("pyrogram").setLevel(logging.WARNING)
 logging.getLogger("PIL").setLevel(logging.WARNING)
-LOG_ = logging.getLogger("core")
+logging.getLogger("pikepdf").setLevel(logging.WARNING)

pdfbot.py

@@ -1,8 +1,13 @@
-from pool import Task, Worker
+from pyrogram.handlers.handler import Handler
+from pool import TaskPool, Worker
 from pyrogram import Client
+from typing import Optional
+from logger import logging
 from creds import my
 import yaml
 
+_LOG = logging.getLogger(__name__)
+
 
 class Pdfbot(Client):
  def __init__(self):
@@ -14,14 +19,29 @@ def __init__(self):
  plugins=dict(root="plugins"),
  )
  self.process_pool = Worker()
- self.task_pool = Task()
+ self.task_pool = TaskPool()
  self.languages: dict = {}
 
+ def new_task(
+ self,
+ _type,
+ chat_id: int,
+ message_id: int,
+ handler: Optional[Handler] = None,
+ ):
+ task = _type(chat_id, message_id) # type: ignore
+ if handler:
+ task.set_handler(handler)
+ self.add_handler(handler)
+ self.task_pool.add_task(chat_id, task)
+ return task
+
  async def start(self):
  with open("locale/en.yaml", "r") as f:
  self.language = yaml.load(f, Loader=yaml.CLoader)
  await super().start()
  await self.process_pool.start()
+ _LOG.info("Client started")
 
  async def stop(self):
  await self.process_pool.stop()

plugins/callback_handler.py

@@ -0,0 +1,92 @@
+# (c) AlenPaulVarghese
+# -*- coding: utf-8 -*-
+
+from tools.general import rotate_image # pylint:disable=import-error
+from pyrogram import filters
+from pdfbot import Pdfbot # pylint:disable=import-error
+from pyrogram.types import (
+ CallbackQuery,
+ InputMediaPhoto,
+ ReplyKeyboardRemove,
+)
+from tools import MakePdf
+import asyncio
+
+
+@Pdfbot.on_callback_query(filters.regex(r"insert|remove|rotate|^rm_task$|^del$"))
+async def callback_handler(client: Pdfbot, callback: CallbackQuery):
+ current_task = client.task_pool.get_task(callback.message.chat.id)
+ if current_task is None:
+ await asyncio.gather(
+ callback.answer("cancelled/timed-out"),
+ callback.message.delete(),
+ )
+ return
+ elif callback.data == "rm_task":
+ client.task_pool.remove_task(callback.message.chat.id)
+ await asyncio.gather(
+ callback.message.reply_text(
+ "**Task** cancelled", reply_markup=ReplyKeyboardRemove()
+ ),
+ callback.message.delete(),
+ )
+ return
+ elif callback.data == "del":
+ await callback.message.delete()
+ return
+
+ file_id = int(callback.data.split(":", 1)[0])
+ file_path = current_task.temp_files.get(file_id)
+ if file_path is None:
+ await asyncio.gather(
+ callback.answer("cancelled/timed-out"),
+ callback.message.delete(),
+ )
+ elif "rotate" in callback.data:
+ await callback.answer("rotating please wait")
+ degree = int(callback.data.split(":", 2)[2])
+ temporary_image = await asyncio.get_event_loop().run_in_executor(
+ None, rotate_image, file_path, degree
+ )
+ await callback.message.edit_media(
+ InputMediaPhoto(temporary_image), reply_markup=callback.message.reply_markup
+ )
+ elif "insert" in callback.data:
+ if file_path is not None:
+ current_task.temp_files.pop(file_id)
+ current_task.proposed_files.append(file_path)
+ await callback.message.delete()
+ if not current_task.quiet:
+ await callback.message.reply_text(
+ text="photo added successfully"
+ if isinstance(current_task, MakePdf)
+ else "pdf added successfully"
+ )
+ elif "remove" in callback.data:
+ current_task.temp_files.pop(file_id)
+ await callback.message.delete()
+
+
+# @Pdfbot.on_message(filters.command(["extract"]) & filters.create(_task_checker))
+# async def extract_handler(client: Pdfbot, message: Message) -> None:
+# new_task = Extractor(
+# client, message.chat.id, message.message_id, " ".join(message.command[1:])
+# )
+# client.task_pool.add_task(message.chat.id, new_task)
+# if not await new_task.parse_input():
+# return
+# status = await message.reply_text("**downloading...**")
+# await new_task.allocate_and_download(message.reply_to_message)
+# await status.edit(
+# "**processing...**",
+# )
+# client.process_pool.new_task(new_task)
+# while new_task.status == 0:
+# await asyncio.sleep(1.2)
+# else:
+# if new_task.status == 2:
+# await message.reply_text(f"**Task failed: `{new_task.error_code}`**")
+# elif new_task.status == 1:
+# await message.reply_document(new_task.output)
+# new_task.__del__()
+# client.task_pool.remove_task(message.chat.id)

plugins/task_handlers/decrypt_handler.py

@@ -0,0 +1,46 @@
+# (c) AlenPaulVarghese
+# -*- coding: utf-8 -*-
+
+from tools import task_checker, slugify
+from pyrogram.types import Message
+from pyrogram import filters
+from tools import Decrypter
+from pdfbot import Pdfbot
+
+
+@Pdfbot.on_message(filters.command(["decrypt"]) & filters.create(task_checker))
+async def encrypt_handler(client: Pdfbot, message: Message) -> None:
+ if message.reply_to_message is None or (
+ message.document and message.document.mime_type != "document/pdf"
+ ):
+ await message.reply("Please reply to a PDF file")
+ return
+ if len(message.command) == 1:
+ await message.reply("**usage:** `/decrypt password custom_file_name`")
+ return
+ status = await message.reply_text("**downloading...**")
+ task = client.new_task(Decrypter, message.chat.id, message.message_id)
+ if len(message.command) >= 3:
+ filename = " ".join(message.command[2:])
+ if len(filename) > 64:
+ await message.reply_text("**WARNING:** file name too long, reducing...")
+ filename = slugify(filename[:64])
+ if filename in {".", "..", ".pdf"}:
+ await message.reply_text(
+ "**WARNING:** invalid filename falling back to default."
+ )
+ else:
+ task.filename = filename
+ input_file = task.cwd / f"{message.reply_to_message.message_id}.pdf"
+ await message.reply_to_message.download(input_file)
+ task.set_configuration(input_file, message.command[1])
+ await status.edit("**processing...**")
+ try:
+ await client.process_pool.new_task(task)
+ await message.reply_document(task.cwd / task.filename)
+ await status.delete()
+ except Exception as e:
+ await status.edit(f"**Task failed: `{e}`**")
+ finally:
+ task.cleanup()
+ client.task_pool.remove_task(message.chat.id)