Merge pull request #181 from aahnik/upgrade-plugins

Upgrade plugins

Dockerfile

@@ -5,6 +5,7 @@ WORKDIR /app
 RUN apt-get update && \
     apt-get install -y --no-install-recommends apt-utils && \
     apt-get upgrade -y && \
+    apt-get install ffmpeg tesseract-ocr -y && \
     apt-get autoclean
 RUN pip install --upgrade poetry
 RUN python -m venv /venv

pyproject.toml

@@ -1,19 +1,12 @@
 [tool.poetry]
 name = "tgcf"
-version = "0.1.35"
+version = "0.2.0"
 description = "The ultimate tool to automate custom telegram message forwarding."
 authors = ["aahnik <[email protected]>"]
 license = "MIT"
 readme = "README.md"
 repository = "https://github.com/aahnik/tgcf"
 documentation = "https://github.com/aahnik/tgcf/wiki"
-packages = [
-    { include = "tgcf"},
-    { include = "tgcf_filter", from = "plugins" },
-    { include = "tgcf_replace", from = "plugins" },
-    { include = "tgcf_watermark", from = "plugins" },
-    { include = "tgcf_ocr", from = "plugins" },
-]
 
 [tool.poetry.dependencies]
 python = "^3.8"
@@ -28,6 +21,8 @@ Pillow = "^8.1.2"
 hachoir = "^3.1.2"
 aiohttp = "^3.7.4"
 tg-login = "^0.0.2"
+"watermark.py" = "^0.0.3"
+pytesseract = "^0.3.7"
 
 [tool.poetry.dev-dependencies]
 pytest = "^5.2"

tgcf/live.py

@@ -2,6 +2,7 @@
 import logging
 
 from telethon import events
+from telethon.tl.custom.message import Message
 
 from tgcf import config, const
 from tgcf.bot import BOT_EVENTS
@@ -10,7 +11,7 @@
 
 existing_hashes = []
 
-_stored = {}
+from typing import Dict, List
 
 
 class EventUid:
@@ -33,6 +34,9 @@ def __hash__(self) -> int:
         return hash(self.__str__())
 
 
+_stored: Dict[EventUid, List[Message]] = {}
+
+
 async def new_message_handler(event):
     """Process new incoming messages."""
     chat_id = event.chat_id
@@ -60,13 +64,13 @@ async def new_message_handler(event):
         if event_uid not in _stored:
             _stored[event_uid] = []
 
-        message = apply_plugins(message)
-        if not message:
+        tm = await apply_plugins(message)
+        if not tm:
             return
         for recipient in to_send_to:
-            fwded_msg = await send_message(event.client, recipient, message)
+            fwded_msg = await send_message(recipient, tm)
             _stored[event_uid].append(fwded_msg)
-
+        tm.clear()
     existing_hashes.append(hash(message.text))
 
 
@@ -83,9 +87,9 @@ async def edited_message_handler(event):
 
     event_uid = EventUid(event)
 
-    message = apply_plugins(message)
+    tm = await apply_plugins(message)
 
-    if not message:
+    if not tm:
         return
 
     fwded_msgs = _stored.get(event_uid)
@@ -96,13 +100,14 @@ async def edited_message_handler(event):
                 await msg.delete()
                 await message.delete()
             else:
-                await msg.edit(message.text)
+                await msg.edit(tm.text)
         return
 
     to_send_to = config.from_to.get(event.chat_id)
 
     for recipient in to_send_to:
-        await send_message(event.client, recipient, message)
+        await send_message(recipient, tm)
+    tm.clear()
 
 
 async def deleted_message_handler(event):

tgcf/past.py

@@ -29,12 +29,13 @@ async def forward_job():
                 if isinstance(message, MessageService):
                     continue
                 try:
-                    message = apply_plugins(message)
-                    if not message:
+                    tm = await apply_plugins(message)
+                    if not tm:
                         continue
 
                     for destination in forward.dest:
-                        await send_message(client, destination, message)
+                        await send_message(destination, tm)
+                    tm.clear()
                     last_id = str(message.id)
                     logging.info(f"forwarding message with id = {last_id}")
                     forward.offset = last_id

tgcf/plugins/__init__.py

@@ -0,0 +1,140 @@
+"""Subpackage of tgcf: plugins.
+
+Contains all the first-party tgcf plugins.
+"""
+
+
+import inspect
+import logging
+from enum import Enum
+from importlib import import_module
+from typing import Any, Dict
+
+from telethon.tl.custom.message import Message
+
+from tgcf.config import CONFIG
+from tgcf.utils import cleanup, stamp
+
+PLUGINS = CONFIG.plugins
+
+
+class FileType(str, Enum):
+    AUDIO = "audio"
+    GIF = "gif"
+    VIDEO = "video"
+    VIDEO_NOTE = "video_note"
+    STICKER = "sticker"
+    CONTACT = "contact"
+    PHOTO = "photo"
+    DOCUMENT = "document"
+    NOFILE = "nofile"
+
+
+class TgcfMessage:
+    def __init__(self, message: Message):
+        self.message = message
+        self.text = self.message.text
+        self.raw_text = self.message.raw_text
+        self.sender_id = self.message.sender_id
+        self.file_type = self.guess_file_type()
+        self.new_file = None
+        self.cleanup = False
+
+    async def get_file(self):
+        """Downloads the file in the message and returns the path where its saved."""
+        self.file = stamp(await self.message.download_media(""), self.sender_id)
+        return self.file
+
+    def guess_file_type(self):
+        for i in FileType:
+            if i == FileType.NOFILE:
+                return i
+            obj = getattr(self.message, i.value)
+            if obj:
+                return i
+
+    def clear(self):
+        if self.new_file and self.cleanup:
+            cleanup(self.new_file)
+            self.new_file = None
+
+
+class TgcfPlugin:
+    id_ = "plugin"
+
+    def __init__(self, data: Dict[Any, Any]):
+        self.data = data
+
+    def modify(self, tm: TgcfMessage):
+        """Modify the message here."""
+        return tm
+
+
+def load_plugins() -> Dict[str, TgcfPlugin]:
+    """Load the plugins specified in config."""
+    _plugins = {}
+    for plugin_id, plugin_data in PLUGINS.items():
+        if not plugin_data:
+            plugin_data = {}
+        plugin_class_name = f"Tgcf{plugin_id.title()}"
+
+        try:
+            plugin_module = import_module("tgcf.plugins." + plugin_id)
+        except ModuleNotFoundError:
+            logging.error(
+                f"{plugin_id} is not a first party plugin. Trying to load from availaible third party plugins."
+            )
+            try:
+                plugin_module_name = f"tgcf_{plugin_id}"
+                plugin_module = import_module(plugin_module_name)
+            except ModuleNotFoundError:
+                logging.error(
+                    f"Module {plugin_module_name} not found. Failed to load plugin {plugin_id}"
+                )
+                continue
+            else:
+                logging.info(
+                    f"Plugin {plugin_id} is successfully loaded from third party plugins"
+                )
+        else:
+            logging.info(f"First party plugin {plugin_id} loaded!")
+        try:
+            plugin_class = getattr(plugin_module, plugin_class_name)
+            if not issubclass(plugin_class, TgcfPlugin):
+                logging.error(
+                    f"Plugin class {plugin_class_name} does not inherit TgcfPlugin"
+                )
+                continue
+            plugin: TgcfPlugin = plugin_class(plugin_data)
+            if not plugin.id_ == plugin_id:
+                logging.error(f"Plugin id for {plugin_id} does not match expected id.")
+                continue
+        except AttributeError:
+            logging.error(f"Found plugin {plugin_id}, but plguin class not found.")
+        else:
+            logging.info(f"Loaded plugin {plugin_id}")
+            _plugins.update({plugin.id_: plugin})
+    return _plugins
+
+
+async def apply_plugins(message: Message) -> TgcfMessage:
+    """Apply all loaded plugins to a message."""
+    tm = TgcfMessage(message)
+
+    for _id, plugin in plugins.items():
+        try:
+            if inspect.iscoroutinefunction(plugin.modify):
+                ntm = await plugin.modify(tm)
+            else:
+                ntm = plugin.modify(tm)
+        except Exception as err:
+            logging.error(f"Failed to apply plugin {_id}. \n {err} ")
+        else:
+            logging.info(f"Applied plugin {_id}")
+            if not ntm:
+                tm.clear()
+                return None
+    return tm
+
+
+plugins = load_plugins()