From 44803cf4a1b6396afa572a89d3e3b458c65b2250 Mon Sep 17 00:00:00 2001
From: Hehesheng <hhslzzyyes@outlook.com>
Date: Tue, 4 Jun 2024 00:01:59 +0800
Subject: [PATCH] feat: supported link convert

---
 backend/MediaCacheManager.py         |  2 +-
 backend/TgFileSystemClient.py        |  8 +++++++-
 backend/TgFileSystemClientManager.py |  5 +++++
 backend/api.py                       | 30 +++++++++++++++++++++++++---
 4 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/backend/MediaCacheManager.py b/backend/MediaCacheManager.py
index 9ac6064..3ce2b50 100644
--- a/backend/MediaCacheManager.py
+++ b/backend/MediaCacheManager.py
@@ -167,7 +167,7 @@ class MediaChunkHolderManager(object):
     def __init__(self) -> None:
         self.chunk_lru = collections.OrderedDict()
         self.disk_chunk_cache = diskcache.Cache(
-            f"{os.path.dirname(__file__)}/cache_media", size_limit=MediaChunkHolderManager.MAX_CACHE_SIZE
+            f"{os.path.dirname(__file__)}/cache_media", size_limit=MediaChunkHolderManager.MAX_CACHE_SIZE * 2
         )
         self._restore_cache()
 
diff --git a/backend/TgFileSystemClient.py b/backend/TgFileSystemClient.py
index c2e6bb1..93ef2f2 100644
--- a/backend/TgFileSystemClient.py
+++ b/backend/TgFileSystemClient.py
@@ -37,6 +37,7 @@ class TgFileSystemClient(object):
     qr_login: QRLogin | None = None
     login_task: asyncio.Task | None = None
     # rsa key
+    sign: str
     public_key: rsa.PublicKey
     private_key: rsa.PrivateKey
     # task should: (task_id, callabledFunc)
@@ -69,6 +70,7 @@ class TgFileSystemClient(object):
             (client_param for client_param in param.clients if client_param.token == session_name),
             configParse.TgToFileSystemParameter.ClientConfigPatameter(),
         )
+        self.sign = self.client_param.token
         self.task_queue = asyncio.Queue()
         self.client = TelegramClient(
             f"{os.path.dirname(__file__)}/db/{self.session_name}.session",
@@ -210,7 +212,7 @@ class TgFileSystemClient(object):
             logger.info(f"{chat_id} quit cache task.")
 
     @_acheck_before_call
-    async def get_message(self, chat_id: int, msg_id: int) -> types.Message:
+    async def get_message(self, chat_id: int | str, msg_id: int) -> types.Message:
         msg = await self.client.get_messages(chat_id, ids=msg_id)
         return msg
 
@@ -245,6 +247,10 @@ class TgFileSystemClient(object):
             offset = first_id + offset
         return offset
 
+    @_acheck_before_call
+    async def get_entity(self, chat_id_or_name) -> hints.Entity:
+        return await self.client.get_entity(chat_id_or_name)
+
     @_acheck_before_call
     async def get_messages(self, chat_id: int, limit: int = 10, offset: int = 0) -> hints.TotalList:
         offset = await self._get_offset_msg_id(chat_id, offset)
diff --git a/backend/TgFileSystemClientManager.py b/backend/TgFileSystemClientManager.py
index 71cbfba..b98f901 100644
--- a/backend/TgFileSystemClientManager.py
+++ b/backend/TgFileSystemClientManager.py
@@ -86,6 +86,11 @@ class TgFileSystemClientManager(object):
         client = self.clients.get(client_id)
         return client
 
+    def get_first_client(self) -> TgFileSystemClient:
+        for client in self.clients.values():
+            return client
+        return None
+
     async def get_client_force(self, client_id: str) -> TgFileSystemClient:
         client = self.get_client(client_id)
         if client is None:
diff --git a/backend/api.py b/backend/api.py
index 16e26a3..df99168 100644
--- a/backend/api.py
+++ b/backend/api.py
@@ -9,7 +9,7 @@ from fastapi import FastAPI, status, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import Response, StreamingResponse
 from contextlib import asynccontextmanager
-from telethon import types, hints
+from telethon import types, hints, utils
 from pydantic import BaseModel
 
 import configParse
@@ -199,8 +199,32 @@ async def get_tg_file_client_status(request: Request):
 
 @app.get("/tg/api/v1/client/link_convert")
 @apiutils.atimeit
-async def convert_tg_msg_link_media_stream(link: str, sign: str):
-    raise NotImplementedError
+async def convert_tg_msg_link_media_stream(link: str):
+    try:
+        link_slice = link.split("/")
+        if len(link_slice) < 5:
+            raise RuntimeError("link format invalid")
+        chat_id_or_name, msg_id = link_slice[-2:]
+        is_msg_id = msg_id.isascii() and msg_id.isdecimal()
+        if not is_msg_id:
+            raise RuntimeError("message id invalid")
+        msg_id = int(msg_id)
+        is_chat_name = chat_id_or_name.isascii() and not chat_id_or_name.isdecimal()
+        is_chat_id = chat_id_or_name.isascii() and chat_id_or_name.isdecimal()
+        if not is_chat_name and not is_chat_id:
+            raise RuntimeError("chat id invalid")
+        client = clients_mgr.get_first_client()
+        if client is None:
+            raise RuntimeError("client not ready, login first pls.")
+        if is_chat_id:
+            chat_id_or_name = int(chat_id_or_name)
+        msg = await client.get_message(chat_id_or_name, msg_id)
+        file_name = apiutils.get_message_media_name(msg)
+        url = f"{param.base.exposed_url}/tg/api/v1/file/get/{utils.get_peer_id(msg.peer_id)}/{msg.id}/{file_name}?sign={client.sign}"
+        return Response(json.dumps({"url": url}), status_code=status.HTTP_200_OK)
+    except Exception as err:
+        logger.error(f"{err=}")
+        return Response(json.dumps({"detail": "link invalid", "err": f"{err}"}), status_code=status.HTTP_404_NOT_FOUND)
 
 
 @app.get("/tg/api/v1/client/profile_photo")