Merge pull request #121 from aahnik/add-more-filters

Add more filters

plugins/tgcf_filter/__init__.py

@@ -1,4 +1,5 @@
 import logging
+from enum import Enum
 from typing import List, Optional
 
 from pydantic import BaseModel
@@ -9,8 +10,27 @@ class FilterList(BaseModel):
     whitelist: Optional[List[str]] = []
 
 
+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 FilesFilterList(BaseModel):
+    blacklist: Optional[List[FileType]] = []
+    whitelist: Optional[List[FileType]] = []
+
+
 class Filters(BaseModel):
-    text: Optional[FilterList]
+    users: Optional[FilterList] = FilterList()
+    files: Optional[FilesFilterList] = FilesFilterList()
+    text: Optional[FilterList] = FilterList()
 
 
 class TgcfFilter:
@@ -22,14 +42,17 @@ def __init__(self, data):
         logging.info(self.filters)
 
     def modify(self, message):
-        msg_text = message.text
-        if not msg_text and self.filters.text.whitelist == []:
-            return message
 
-        if self.text_safe(msg_text, self.filters.text):
-            return message
+        if self.users_safe(message):
+            if self.files_safe(message):
+                if self.text_safe(message):
+                    return message
 
-    def text_safe(self, text, flist: FilterList):
+    def text_safe(self, message):
+        flist = self.filters.text
+        text = message.text
+        if not text and flist.whitelist == []:
+            return True
         for forbidden in flist.blacklist:
             if forbidden in text:
                 return False
@@ -38,3 +61,44 @@ def text_safe(self, text, flist: FilterList):
         for allowed in flist.whitelist:
             if allowed in text:
                 return True
+
+    def users_safe(self, message):
+        flist = self.filters.users
+        sender = str(message.sender_id)
+        logging.info(f"M message from sender id {sender}")
+        if sender in flist.blacklist:
+            return False
+        if not flist.whitelist:
+            return True
+        if sender in flist.whitelist:
+            return True
+
+    def files_safe(self, message):
+        flist = self.filters.files
+
+        def file_type(message):
+            file_types = [
+                FileType.AUDIO,
+                FileType.GIF,
+                FileType.VIDEO,
+                FileType.VIDEO_NOTE,
+                FileType.STICKER,
+                FileType.CONTACT,
+                FileType.PHOTO,
+                FileType.DOCUMENT,
+            ]
+            for _type in file_types:
+                obj = getattr(message, _type)
+                if obj:
+                    return _type
+
+        fl_type = file_type(message)
+        if not fl_type:
+            fl_type = FileType.NOFILE
+        print(fl_type)
+        if fl_type in flist.blacklist:
+            return False
+        if not flist.whitelist:
+            return True
+        if fl_type in flist.whitelist:
+            return True

plugins/tgcf_ocr/__init__.py

@@ -1,4 +1,3 @@
-
 class TgcfOcr:
     id = "ocr"
 

plugins/tgcf_watermark/__init__.py

@@ -1,4 +1,3 @@
-
 class TgcfWatermark:
     id = "watermark"
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tgcf"
-version = "0.1.26"
+version = "0.1.27"
 description = "The ultimate tool to automate telegram message forwarding."
 authors = ["aahnik <[email protected]>"]
 license = "MIT"

tests/test_config.py

@@ -1,5 +1,5 @@
-import os
 import logging
+import os
 
 logging.basicConfig(level=logging.INFO)
 

tgcf/config.py

@@ -57,7 +57,7 @@ def detect_config_type() -> int:
 CONFIG_TYPE = detect_config_type()
 
 
-def read_config_file()->Config:
+def read_config_file() -> Config:
     if CONFIG_TYPE == 1:
         with open(CONFIG_FILE_NAME) as file:
             config_dict = yaml.full_load(file)