From b0e03ce6e8f72647f0f5ce39e657cca806b836d8 Mon Sep 17 00:00:00 2001 From: vvzvlad Date: Mon, 3 Feb 2025 02:28:58 +0300 Subject: [PATCH] Enhance media retrieval with file unique ID and improved error handling --- api_server.py | 66 +++++++++++++++++++++++++++++++++++----------- post_parser.py | 10 +++---- rss_generator.py | 6 +---- telegram_client.py | 1 - 4 files changed, 57 insertions(+), 26 deletions(-) diff --git a/api_server.py b/api_server.py index 1f963a0..c8eaa8c 100644 --- a/api_server.py +++ b/api_server.py @@ -27,6 +27,44 @@ if __name__ == "__main__": import uvicorn uvicorn.run("api_server:app", host=Config["api_host"], port=Config["api_port"], reload=True) + +async def find_file_id_in_message(message, file_unique_id: str): + """Find file_id by checking all possible media types in message""" + if message.photo and message.photo.file_unique_id == file_unique_id: + return message.photo.file_id + elif message.video and message.video.file_unique_id == file_unique_id: + return message.video.file_id + elif message.animation and message.animation.file_unique_id == file_unique_id: + return message.animation.file_id + elif message.video_note and message.video_note.file_unique_id == file_unique_id: + return message.video_note.file_id + elif message.audio and message.audio.file_unique_id == file_unique_id: + return message.audio.file_id + elif message.voice and message.voice.file_unique_id == file_unique_id: + return message.voice.file_id + elif message.sticker and message.sticker.file_unique_id == file_unique_id: + return message.sticker.file_id + elif message.web_page and message.web_page.photo and message.web_page.photo.file_unique_id == file_unique_id: + return message.web_page.photo.file_id + elif message.document and message.document.file_unique_id == file_unique_id: + return message.document.file_id + return None + +async def prepare_file_response(file_path: str, background_tasks: BackgroundTasks): + """Prepare file response with proper headers and cleanup task""" + if not os.path.exists(file_path): + raise HTTPException(status_code=404, detail="Temporary file not found") + + media_type, _ = mimetypes.guess_type(file_path) + if not media_type: + media_type = "application/octet-stream" + + background_tasks.add_task(lambda: os.remove(file_path) if os.path.exists(file_path) else None) + + headers = {"Content-Disposition": f"inline; filename={os.path.basename(file_path)}"} + return FileResponse(path=file_path, media_type=media_type, headers=headers) + + @app.get("/html/{channel}/{post_id}", response_class=HTMLResponse) @app.get("/post/html/{channel}/{post_id}", response_class=HTMLResponse) async def get_post_html(channel: str, post_id: int): @@ -74,27 +112,25 @@ async def health_check(): logger.error(error_message) raise HTTPException(status_code=500, detail=error_message) from e -@app.get("/media/{file_id}") -async def get_media(file_id: str, background_tasks: BackgroundTasks): +@app.get("/media/{channel}/{post_id}/{file_unique_id}") +async def get_media(channel: str, post_id: int, file_unique_id: str, background_tasks: BackgroundTasks): try: - file_path = await client.download_media_file(file_id) - logger.info(f"Downloaded media file {file_id} to {file_path}") - if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="Temporary file not found") - - media_type, _ = mimetypes.guess_type(file_path) # Determine media type from file extension - if not media_type: media_type = "application/octet-stream" - - background_tasks.add_task(lambda: os.remove(file_path) if os.path.exists(file_path) else None) - - headers = { "Content-Disposition": f"inline; filename={os.path.basename(file_path)}" } - return FileResponse( path=file_path, media_type=media_type, headers=headers) + message = await client.client.get_messages(channel, post_id) + file_id = await find_file_id_in_message(message, file_unique_id) + + if not file_id: + logger.error(f"Media with file_unique_id {file_unique_id} not found in message {post_id}") + raise HTTPException(status_code=404, detail="File not found in message") + file_path = await client.client.download_media(file_id) + logger.info(f"Downloaded media file {file_unique_id} to {file_path}") + return await prepare_file_response(file_path, background_tasks) except HTTPException: raise except errors.RPCError as e: - logger.error(f"Media request RPC error {file_id}: {type(e).__name__} - {str(e)}") + logger.error(f"Media request RPC error for file_unique_id {file_unique_id} in message {post_id}: {type(e).__name__} - {str(e)}") raise HTTPException(status_code=404, detail="File not found in Telegram") from e except Exception as e: - error_message = f"Failed to get media for file_id {file_id}: {str(e)}" + error_message = f"Failed to get media for file_unique_id {file_unique_id} in message {post_id}: {str(e)}" logger.error(error_message) raise HTTPException(status_code=500, detail=error_message) from e diff --git a/post_parser.py b/post_parser.py index dacb82d..a8accdb 100644 --- a/post_parser.py +++ b/post_parser.py @@ -24,7 +24,7 @@ logger = logging.getLogger(__name__) #http://127.0.0.1:8000/post/html/DragorWW_space/61 — links #http://127.0.0.1:8000/post/html/theyforcedme/3577 - video note #http://127.0.0.1:8000/post/html/theyforcedme/3572 - audio -#http://127.0.0.1:8000/post/html/theyforcedme/3558 audio-note +#http://127.0.0.1:8000/post/html/theyforcedme/3558 - audio-note #http://127.0.0.1:8000/html/vvzvlad_lytdybr/426 - sticker #http://127.0.0.1:8000/html/wrkshprn/634 — links without @@ -65,7 +65,7 @@ class PostParser: raise def _format_json(self, message: Message, naked: bool = False) -> Dict[Any, Any]: - html_content = self._format_html(message, naked=naked) + html_content = self._format_html(message) result = { 'channel': message.chat.username, 'message_id': message.id, @@ -173,7 +173,7 @@ class PostParser: if webpage := getattr(message, "web_page", None): # Web page preview - if webpage_html := self._format_webpage(webpage): + if webpage_html := self._format_webpage(webpage, message): html_content.append(webpage_html) if text: # Message text @@ -218,12 +218,12 @@ class PostParser: logger.error(f"url_processing_error: error {str(e)}") return text - def _format_webpage(self, webpage) -> Union[str, None]: + def _format_webpage(self, webpage, message) -> Union[str, None]: base_url = Config['pyrogram_bridge_url'] try: if photo := getattr(webpage, "photo", None): if file_unique_id := getattr(photo, "file_unique_id", None): - url = f"{base_url}/media/{webpage.chat.username}/{webpage.id}/{file_unique_id}" + url = f"{base_url}/media/{message.chat.username}/{message.id}/{file_unique_id}" return ( f'
' f'' diff --git a/rss_generator.py b/rss_generator.py index c3f9d9c..985c52c 100644 --- a/rss_generator.py +++ b/rss_generator.py @@ -38,7 +38,6 @@ async def generate_channel_rss(channel: str, post_parser: Optional[PostParser] = try: channel_info = await post_parser.client.get_chat(channel) channel_title = channel_info.title or f"Telegram: {channel}" - channel_icon = getattr(channel_info.photo, 'small_file_id', None) if channel_info.photo else None except Exception as e: # raise error if channel not found if "USERNAME_INVALID" in str(e): return create_error_feed(channel, base_url) @@ -53,9 +52,6 @@ async def generate_channel_rss(channel: str, post_parser: Optional[PostParser] = fg.language('ru') #fg.dc.dc_creator(f"@{channel}") - if channel_icon: - fg.logo(f"{base_url}/media/{channel_icon}") - fg.icon(f"{base_url}/media/{channel_icon}") fg.id(f"{base_url}/rss/{channel}") @@ -87,7 +83,7 @@ async def generate_channel_rss(channel: str, post_parser: Optional[PostParser] = continue # Merge media groups - for group_id, group_posts in media_groups.items(): + for _group_id, group_posts in media_groups.items(): if not group_posts: continue diff --git a/telegram_client.py b/telegram_client.py index 970f311..156fb4f 100644 --- a/telegram_client.py +++ b/telegram_client.py @@ -2,7 +2,6 @@ import logging from pyrogram import Client from config import get_settings - logger = logging.getLogger(__name__) settings = get_settings()