Refactor file download logic: enhance filename validation, ensure uniqueness, and update message retrieval methods

Author: Keniis0712

src/db.py

@@ -41,7 +41,6 @@ class Group(Base):
     group_id = Column(Integer, nullable=False)
     name = Column(String, nullable=False)
     last_update = Column(DateTime, nullable=False, default=datetime.datetime.now(datetime.timezone.utc))
-    latest_msgid = Column(Integer, nullable=False)
 
 
 class Message(Base):
@@ -139,8 +138,8 @@ def update_user_session(user_id: int, new_session: str) -> None:
 def add_group(server_id: int, group_id: int, name: str) -> None:
     with SessionLocal() as session:
         group = Group(server_id=server_id, group_id=group_id, name=name,
-                      last_update=datetime.datetime.now(datetime.timezone.utc),
-                      latest_msgid=0)
+                      last_update=datetime.datetime.now(datetime.timezone.utc)
+                      )
         session.add(group)
         session.commit()
 
@@ -154,14 +153,6 @@ def update_group_name(group_id: int, server_id: int, new_name: str) -> None:
             session.commit()
 
 
-def update_group_msgid(group_id: int, server_id: int, msgid: int) -> None:
-    with SessionLocal() as session:
-        group = session.query(Group).filter_by(group_id=group_id, server_id=server_id).first()
-        if group:
-            group.latest_msgid = msgid
-            session.commit()
-
-
 def get_group_msgid(group_id: int, server_id: int, latest: bool = True) -> Optional[int]:
     if latest:
         data = get_messages(server_id, group_id, MAX_MSGID, False, 1)
@@ -175,7 +166,6 @@ def get_group_msgid(group_id: int, server_id: int, latest: bool = True) -> Optio
         return MAX_MSGID
 
 
-
 async def get_group_name(server_id: int, user: StealthIM.User, group_id: int, force_flush=False) -> StealthIM.group.GroupPublicInfoResult:
     with SessionLocal() as session:
         group = session.query(Group).filter_by(group_id=group_id, server_id=server_id).first()

src/screens/chat.py

@@ -1,6 +1,7 @@
 import asyncio
 import datetime
 import math
+import os
 from typing import Optional, cast
 
 import platformdirs
@@ -283,7 +284,22 @@ async def action_download(self):
 
         download_path = platformdirs.user_downloads_path()
         hashes = [msg.hash for msg in self.selected]
-        filenames = [msg.text for msg in files]
+        filenames = []
+        for msg in files:
+            filename = msg.text
+            # Remove invalid characters
+            filename = "".join(c for c in filename if c not in r'\/:*?"<>|')
+            if not filename:
+                filename = msg.hash
+            # Ensure the filename is unique
+            file_path = os.path.join(download_path, filename)
+            if os.path.exists(file_path):
+                base, ext = os.path.splitext(filename)
+                i = 1
+                while os.path.exists(os.path.join(download_path, f"{base}({i}){ext}")):
+                    i += 1
+                filename = f"{base}({i}){ext}"
+            filenames.append(filename)
         await self.app.data.group.download_files(hashes, filenames, download_path)
 
     @staticmethod
@@ -298,7 +314,7 @@ class ChatScreen(Screen):
     SCREEN_NAME = "Chat"
     CSS_PATH = "../../styles/chat.tcss"
 
-    LIMIT = 100
+    LIMIT = 10
 
     BINDINGS = [("ctrl+s", "select_msg", "Select message")]
 
@@ -396,7 +412,7 @@ async def on_change_group(self, event: ListView.Selected) -> None:
         else:
             # A new group, we only get the newest LIMIT messages
             # from_id=0, old_to_new=False means pull the latest messages
-            gen = self.group.receive_text(from_id=0, old_to_new=False, sync=False, limit=self.LIMIT)
+            gen = self.group.receive_latest_text(limit=self.LIMIT)
             msgs = [x async for x in gen][::-1]
             for msg in msgs:
                 message = db.add_message(
@@ -442,8 +458,8 @@ async def on_no_more_message(self, _event: TopDetectingScroll.ScrolledToTop):
         else:
             # There's no more messages in the database, try to pull from server
             oldest_msgid = db.get_group_msgid(self.group.group_id, self.app.data.server_db.id, False)
-            gen = self.group.receive_text(from_id=oldest_msgid, old_to_new=False, sync=False, limit=self.LIMIT)
-            msgs = [x async for x in gen][::-1]
+            gen = self.group.receive_text(from_id=oldest_msgid, sync=False, limit=self.LIMIT)
+            msgs = [x async for x in gen]
             self.log(msgs)
             for msg in msgs:
                 message = db.add_message(
@@ -605,23 +621,21 @@ async def get_messages(self, messages: VerticalScroll) -> None:
         group_id = self.group.group_id
 
         while True:
-            latest_msgid = db.get_group_msgid(group_id, server_id)
             try:
-                gen = self.group.receive_text(from_id=latest_msgid)
+                gen = self.group.receive_new_text(limit=self.LIMIT)
                 async for message in gen:
                     msg = db.add_message(
                         server_id, group_id, message.type.value,
                         message.msg.replace("\n", "\n\n"),
                         datetime.datetime.fromtimestamp(int(message.time)), message.username,
                         message.msgid, message.hash
                     )
-                    db.update_group_msgid(group_id, server_id, message.msgid)
 
-                    if message.type != MessageType.Recall:
-                        await self.add_message(messages, self.build_msg_from_db(msg))
-                    else:
-                        db.recall_message(server_id, group_id, message.msgid)
-                        await self.recall_message(messages, message.msgid)
+                    # if message.type != MessageType.Recall:
+                    await self.add_message(messages, self.build_msg_from_db(msg))
+                    # else:
+                    #     db.recall_message(server_id, group_id, message.msgid)
+                    #     await self.recall_message(messages, message.msgid)
             except RuntimeError:
                 pass
             except asyncio.CancelledError:

src/screens/widgets.py

@@ -109,32 +109,6 @@ def on_changed(self, _event) -> None:
         self.selected = None
 
 
-# noinspection PyProtectedMember
-def add_global_hook(screen):
-    if not hasattr(screen, "_popup_callback_click"):
-        screen._popups = []
-
-        # noinspection PyUnresolvedReferences,PyProtectedMember
-        async def callback(_self, event: Click) -> None:
-            for obj in screen._popups:
-                await obj.on_global_click(event)
-
-        screen._popup_callback_click = callback
-        screen._decorated_handlers.setdefault(Click, [])
-        screen._decorated_handlers[Click].append((callback, None))
-    if not hasattr(screen, "_popup_callback_key"):
-        screen._popups = []
-
-        # noinspection PyUnresolvedReferences,PyProtectedMember
-        async def callback(_self, event: Key) -> None:
-            for obj in screen._popups:
-                await obj.on_global_key(event)
-
-        screen._popup_callback_key = callback
-        screen._decorated_handlers.setdefault(Key, [])
-        screen._decorated_handlers[Key].append((callback, None))
-
-
 class AbstractPopup(Widget):
     def __init__(self, id):
         super().__init__(id=id)
@@ -165,7 +139,7 @@ async def callback(_self, event: Key) -> None:
             screen._decorated_handlers[Key].append((callback, None))
 
     def on_mount(self) -> None:
-        add_global_hook(self.app.screen)
+        self.add_global_hook(self.app.screen)
         # noinspection PyProtectedMember,PyUnresolvedReferences
         self.app.screen._popups.append(self)