Add download and open files support

README.md

@@ -39,5 +39,6 @@ pip3 install python-telegram
 
 Usefull links to help develop this project.
 
-- https://github.com/alexander-akhmetov/python-telegram
-- https://github.com/tdlib/td/tree/master/example#python
+- tdlib python wrapper: https://github.com/alexander-akhmetov/python-telegram
+- tdlib official repo: https://github.com/tdlib/td/tree/master/example#python
+- Build tdlib instructions: https://tdlib.github.io/td/build.html?language=Python

tg/controllers/__init__.py

@@ -1,21 +1,17 @@
 import logging
 import os
 import threading
+from tempfile import NamedTemporaryFile
 
 from telegram.client import Telegram
 
-from utils import notify
+from tg.utils import notify, handle_exception, suspend
 from tg.models import Model
 from tg.views import View
+from tg.msg import MsgProxy
 
 log = logging.getLogger(__name__)
 
-SUPPORTED_MSG_TYPES = (
-    "updateNewMessage",
-    "updateChatLastMessage",
-    "updateMessageSendSucceeded",
-)
-
 
 class Controller:
     """
@@ -25,13 +21,17 @@ class Controller:
     # View is terminal vindow
     """
 
-    tg: Telegram
-
     def __init__(self, model: Model, view: View, tg: Telegram) -> None:
         self.model = model
         self.view = view
-        self.tg = tg
         self.lock = threading.Lock()
+        self.tg = tg
+        self.handlers = {
+            "updateNewMessage": self.update_new_msg,
+            "updateChatLastMessage": self.update_chat_last_msg,
+            "updateMessageSendSucceeded": self.update_msg_send_succeeded,
+            "updateFile": self.update_file,
+        }
 
     def run(self) -> None:
         try:
@@ -39,6 +39,33 @@ class Controller:
         except Exception:
             log.exception("Error happened in main loop")
 
+    def download_current_file(self):
+        msg = MsgProxy(self.model.current_msg())
+        log.debug("Downloading msg: %s", msg.msg)
+        if msg.file_id:
+            log.info("Downloading file: file_id=%s", msg.file_id)
+            self.tg.download_file(file_id=msg.file_id)
+            log.info("Downloaded: file_id=%s", msg.file_id)
+
+    def open_current_msg(self):
+        msg = MsgProxy(self.model.current_msg())
+        log.info("Open msg: %s", msg.msg)
+        if msg.is_text:
+            text = msg['content']['text']["text"]
+            with NamedTemporaryFile('w') as f:
+                f.write(text)
+                f.flush()
+                with suspend(self.view) as s:
+                    s.call(["less", f.name])
+            return
+
+        path = msg.local_path
+        if path:
+            # handle with mimetype and mailcap
+            # if multiple entries in mailcap, open fzf to choose
+            with suspend(self.view) as s:
+                s.call(["open", path])
+
     def handle_msgs(self) -> str:
         # set width to 0.25, move window to left
         # refresh everything
@@ -78,6 +105,11 @@ class Controller:
             elif keys == "dd":
                 if self.model.delete_msg():
                     self.refresh_msgs()
+            elif keys == "D":
+                self.download_current_file()
+
+            elif keys == "l":
+                self.open_current_msg()
 
             elif keys == "/":
                 # search
@@ -163,41 +195,40 @@ class Controller:
         msgs = self.model.fetch_msgs(limit=self.view.msgs.h)
         self.view.draw_msgs(self.model.get_current_chat_msg(), msgs)
 
-    def update_handler(self, update):
-        try:
-            _type = update["@type"]
-            log.info("===Received %s type: %s", _type, update)
-            if _type == "updateNewMessage":
-                # with self.lock:
-                chat_id = update["message"]["chat_id"]
-                self.model.msgs.add_message(chat_id, update["message"])
-                # msgs = self.model.get_current_chat_msg()
-                self.refresh_msgs()
-                if not update.get("disable_notification"):
-                    if update["message"]["content"] == "text":
-                        notify(update["message"]["content"]["text"]["text"])
-            elif _type == "updateChatLastMessage":
-                log.info("Proccessing updateChatLastMessage")
-                chat_id = update["chat_id"]
-                message = update["last_message"]
-                self.model.chats.update_last_message(chat_id, message)
-                self.refresh_chats()
-            elif _type == "updateMessageSendSucceeded":
-                chat_id = update["message"]["chat_id"]
-                msg_id = update["old_message_id"]
-                self.model.msgs.add_message(chat_id, update["message"])
-                self.model.msgs.remove_message(chat_id, msg_id)
-                self.refresh_msgs()
-        except Exception:
-            log.exception("Error happened in update_handler")
-        # message_content = update['message']['content'].get('text', {})
-        # we need this because of different message types: photos, files, etc.
-        # message_text = message_content.get('text', '').lower()
+    @handle_exception
+    def update_new_msg(self, update):
+        chat_id = update["message"]["chat_id"]
+        self.model.msgs.add_message(chat_id, update["message"])
+        # msgs = self.model.get_current_chat_msg()
+        self.refresh_msgs()
+        if not update.get("disable_notification"):
+            if update["message"]["content"] == "text":
+                notify(update["message"]["content"]["text"]["text"])
 
-        # if message_text == 'ping':
-        #     chat_id = update['message']['chat_id']
-        #     # print(f'Ping has been received from {chat_id}')
-        #     self.tg.send_message(
-        #         chat_id=chat_id,
-        #         text='pong',
-        #     )
+    @handle_exception
+    def update_chat_last_msg(self, update):
+        log.info("Proccessing updateChatLastMessage")
+        chat_id = update["chat_id"]
+        message = update["last_message"]
+        self.model.chats.update_last_message(chat_id, message)
+        self.refresh_chats()
+
+    @handle_exception
+    def update_msg_send_succeeded(self, update):
+        chat_id = update["message"]["chat_id"]
+        msg_id = update["old_message_id"]
+        self.model.msgs.add_message(chat_id, update["message"])
+        self.model.msgs.remove_message(chat_id, msg_id)
+        self.refresh_msgs()
+
+    @handle_exception
+    def update_file(self, update):
+        log.info("====: %s", update)
+        file_id = update['file']['id']
+        local = update['file']['local']
+        for msg in map(MsgProxy, self.model.get_current_chat_msgs()):
+            log.info("____: %s, %s", msg.file_id, file_id)
+            if msg.file_id == file_id:
+                msg.local = local
+                self.refresh_msgs()
+                break

tg/main.py

@@ -7,12 +7,12 @@ from functools import partial
 
 from telegram.client import Telegram
 
-from tg.controllers import Controller, SUPPORTED_MSG_TYPES
+from tg.controllers import Controller
 from tg.models import Model
 from tg.views import View
 
 logging.basicConfig(
-    level=logging.DEBUG,
+    level=os.getenv("LOG_LEVEL", "DEBUG"),
     format="%(asctime)s %(levelname)s %(message)s",
     handlers=[
         logging.handlers.RotatingFileHandler(
@@ -34,30 +34,45 @@ def run(tg: Telegram, stdscr: window) -> None:
     view = View(stdscr)
     model = Model(tg)
     controller = Controller(model, view, tg)
-    for msg_type in SUPPORTED_MSG_TYPES:
-        tg.add_update_handler(msg_type, controller.update_handler)
+    for msg_type, handler in controller.handlers.items():
+        tg.add_update_handler(msg_type, handler)
 
     t = threading.Thread(target=controller.run,)
     t.start()
     t.join()
 
 
+class TelegramApi(Telegram):
+    def download_file(
+        self, file_id, priority=16, offset=0, limit=0, synchronous=False,
+    ):
+        result = self.call_method(
+            "downloadFile",
+            params=dict(
+                file_id=file_id,
+                priority=priority,
+                offset=offset,
+                limit=limit,
+                synchronous=synchronous,
+            ),
+            block=False,
+        )
+        result.wait()
+
+
 def main():
     log.debug("#" * 64)
-    tg = Telegram(
+    tg = TelegramApi(
         api_id=API_ID,
         api_hash=API_HASH,
         phone=PHONE,
         database_encryption_key="changeme1234",
+        files_directory=os.path.expanduser("~/.cache/tg/"),
+        tdlib_verbosity=0,
+        # TODO: add in config
+        # library_path="/usr/local/Cellar/tdlib/1.6.0/lib/libtdjson.dylib",
     )
     tg.login()
-
-    # model = Model(tg)
-    # print(model.get_me())
-    # print(model.get_user(246785877))
-    # print(model.chats.get_chat(77769955))
-    # return
-
     wrapper(partial(run, tg))
 
 

tg/models/__init__.py

@@ -2,6 +2,7 @@ import logging
 from collections import defaultdict
 from telegram.client import Telegram
 from typing import Any, Dict, List, Union, Set, Optional
+from tg.msg import MsgProxy
 
 log = logging.getLogger(__name__)
 
@@ -25,12 +26,25 @@ class Model:
             return None
         return self.msgs.current_msgs[chat_id]
 
+    def get_current_chat_msgs(self) -> Optional[int]:
+        chat_id = self.chats.id_by_index(self.current_chat)
+        if chat_id is None:
+            return None
+        return self.msgs.msgs[chat_id]
+
     def fetch_msgs(self, offset: int = 0, limit: int = 10) -> Any:
         chat_id = self.chats.id_by_index(self.current_chat)
         if chat_id is None:
             return []
         return self.msgs.fetch_msgs(chat_id, offset=offset, limit=limit)
 
+    def current_msg(self):
+        chat_id = self.chats.id_by_index(self.current_chat)
+        if chat_id is None:
+            return []
+        current_msg = self.msgs.current_msgs[chat_id]
+        return self.msgs.msgs[chat_id][current_msg]
+
     def jump_bottom(self):
         chat_id = self.chats.id_by_index(self.current_chat)
         return self.msgs.jump_bottom(chat_id)

tg/msg.py

@@ -0,0 +1,106 @@
+import logging
+
+log = logging.getLogger(__name__)
+
+
+class MsgProxy:
+
+    fields_mapping = {
+        'messageDocument': ("document", "document"),
+        'messageVoiceNote': ("voice_note", "voice"),
+        'messageText': ("text", "text"),
+        "messagePhoto": ("photo", "sizes", 0, "photo"),
+        'messageAudio': ("audio", "audio"),
+        'messageVideo': ('video', "video"),
+        'messageVideoNote': ("video_note", "video"),
+    }
+
+    types = {
+        'messageDocument': 'document',
+        'messageVoiceNote': 'voice',
+        'messageText': 'text',
+        "messagePhoto": 'photo',
+        'messageAudio': 'audio',
+        'messageVideo': 'video',
+        'messageVideoNote': 'recording',
+    }
+
+    @classmethod
+    def get_doc(cls, msg, deep=10):
+        doc = msg['content']
+        _type = doc['@type']
+        fields = cls.fields_mapping.get(_type)
+        if fields is None:
+            log.error("msg type not supported: %s", _type)
+            return {}
+        for field in fields[:deep]:
+            if isinstance(field, int):
+                doc = doc[field]
+            else:
+                doc = doc.get(field)
+            if doc is None:
+                return {}
+        return doc
+
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __getitem__(self, key):
+        return self.msg[key]
+
+    def __setitem__(self, key, value):
+        self.msg[key] = value
+
+    @property
+    def type(self):
+        return self.types.get(self.msg['content']['@type'])
+
+    @property
+    def size(self):
+        doc = self.get_doc(self.msg)
+        return doc['size']
+
+    @property
+    def duration(self):
+        if self.type not in ('audio', 'voice'):
+            return None
+        doc = self.get_doc(self.msg, deep=1)
+        return doc['duration']
+
+    @property
+    def file_name(self):
+        if self.type not in ('audio', 'document', 'video'):
+            return None
+        doc = self.get_doc(self.msg, deep=1)
+        return doc['file_name']
+
+    @property
+    def file_id(self):
+        if self.type not in ('audio', 'document', 'photo', 'video', 'recording'):
+            return None
+        doc = self.get_doc(self.msg)
+        return doc['id']
+
+    @property
+    def local_path(self):
+        doc = self.get_doc(self.msg)
+        return doc['local']['path']
+
+    @property
+    def local(self):
+        doc = self.get_doc(self.msg)
+        return doc['local']
+
+    @local.setter
+    def local(self, value):
+        doc = self.get_doc(self.msg)
+        doc['local'] = value
+
+    @property
+    def is_text(self):
+        return self.msg['content']['@type'] == 'messageText'
+
+    @property
+    def is_downloaded(self):
+        doc = self.get_doc(self.msg)
+        return doc['local']['is_downloading_completed']

tg/utils.py

@@ -1,6 +1,9 @@
 import logging
 import os
 from typing import Optional
+import subprocess
+from functools import wraps
+import curses
 
 log = logging.getLogger(__name__)
 
@@ -23,3 +26,30 @@ def notify(msg, subtitle="New message", title="Telegram"):
 
     log.debug("####: %s", f"{cmd} {icon} {sound} {title} {subtitle} {msg}")
     os.system(f"{cmd} {icon} {sound} {title} {subtitle} {msg}")
+
+
+def handle_exception(fun):
+    @wraps(fun)
+    def wrapper(*args, **kwargs):
+        try:
+            return fun(*args, **kwargs)
+        except Exception:
+            log.exception("Error happened in %s handler", fun.__name__)
+    return wrapper
+
+
+class suspend:
+
+    def __init__(self, view):
+        self.view = view
+
+    def call(self, *args, **kwargs):
+        subprocess.call(*args, **kwargs)
+
+    def __enter__(self):
+        curses.endwin()
+        return self
+
+    def __exit__(self, exc_type, exc_val, tb):
+        # works without it, actually
+        curses.doupdate()

tg/views/__init__.py

@@ -4,8 +4,9 @@ import math
 import re
 from datetime import datetime
 
-from utils import num
-from colors import cyan, blue, white, normal, reverse, magenta, get_color
+from tg.utils import num
+from tg.msg import MsgProxy
+from tg.colors import cyan, blue, white, reverse, magenta, get_color
 from _curses import window
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -302,10 +303,7 @@ def get_last_msg(chat: Dict[str, Any]) -> str:
     if not last_msg:
         return "<No messages yet>"
     content = last_msg["content"]
-    _type = content["@type"]
-    if _type == "messageText":
-        return content["text"]["text"]
-    return f"[{_type}]"
+    return parse_content(content)
 
 
 def get_date(chat: Dict[str, Any]) -> str:
@@ -319,10 +317,37 @@ def get_date(chat: Dict[str, Any]) -> str:
 
 
 def parse_content(content: Dict[str, Any]) -> str:
-    _type = content["@type"]
-    if _type == "messageText":
+    msg = MsgProxy({"content": content})
+    log.info("Parsing: %s", msg.msg)
+    if msg.is_text:
         return content["text"]["text"]
-    return f"[{_type}]"
+
+    if not msg.type:
+        # not implemented
+        _type = content['@type']
+        return f"[{_type}]"
+
+    fields = dict(
+        name=msg.file_name,
+        duration=msg.duration,
+        size=msg.size,
+        download=get_download(msg.local, msg.size)
+    )
+    info = ', '.join(
+        f"{k}={v}" for k, v in fields.items() if v is not None
+    )
+
+    return f"[{msg.type}: {info}]"
+
+
+def get_download(local, size):
+    if local['is_downloading_completed']:
+        return "yes"
+    elif local['is_downloading_active']:
+        d = local['downloaded_size']
+        percent = int(d * 100 / size)
+        return f"{percent}%"
+    return "no"
 
 
 emoji_pattern = re.compile(