From c680bf24930a4fb9af8e781a8a7e190efd92734d Mon Sep 17 00:00:00 2001
From: vvzvlad <git@vvzvlad.xyz>
Date: Fri, 7 Feb 2025 01:30:02 +0300
Subject: [PATCH] Enhance media file caching and HTML sanitization with more
 granular controls

---
 api_server.py  | 47 +++++++++++++++++++++++++++++++----------------
 post_parser.py | 27 +++++++++++++++++++--------
 2 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/api_server.py b/api_server.py
index 4c2ec25..07a248e 100644
--- a/api_server.py
+++ b/api_server.py
@@ -203,20 +203,30 @@ async def remove_old_cached_files(media_files: list, cache_dir: str) -> tuple[li
 
     for file_data in media_files:
         try:
-            file_id = file_data.get('file_id')
-            if not file_id:
+            channel = file_data.get('channel')
+            post_id = file_data.get('post_id')
+            file_unique_id = file_data.get('file_unique_id')
+            
+            if not all([channel, post_id, file_unique_id]):
                 continue
 
             last_access_time = file_data.get('added', 0)
             days_since_access = (current_time - last_access_time) / (24 * 3600)
 
             if days_since_access > 20:
-                channel, post_id_str, file_unique_id = file_id.split('-', 2)
-                cache_path = os.path.join(cache_dir, f"{channel}-{post_id_str}-{file_unique_id}")
+                cache_path = os.path.join(cache_dir, str(channel), str(post_id), file_unique_id)
                 
                 if os.path.exists(cache_path):
                     try:
                         os.remove(cache_path)
+                        # Попробуем удалить пустые родительские директории
+                        post_dir = os.path.dirname(cache_path)
+                        channel_dir = os.path.dirname(post_dir)
+                        if not os.listdir(post_dir):
+                            os.rmdir(post_dir)
+                        if not os.listdir(channel_dir):
+                            os.rmdir(channel_dir)
+                            
                         files_removed += 1
                         logger.info(f"Removed old cached file: {cache_path}, last access {days_since_access:.1f} days ago")
                     except Exception as e:
@@ -232,17 +242,18 @@ async def remove_old_cached_files(media_files: list, cache_dir: str) -> tuple[li
 
     # Additionally, remove temporary files with prefix "temp_" if they are older than a threshold (1 hour)
     temp_threshold = 3600  # 1 hour in seconds
-    for file in os.listdir(cache_dir):
-        file_path = os.path.join(cache_dir, file)
-        if os.path.isfile(file_path) and file.startswith("temp_"):
-            try:
-                file_mod_time = os.path.getmtime(file_path)
-                if time.time() - file_mod_time > temp_threshold:
-                    os.remove(file_path)
-                    files_removed += 1
-                    logger.info(f"Removed temporary file: {file_path}")
-            except Exception as e:
-                logger.error(f"Failed to remove temporary file {file_path}: {str(e)}")
+    for root, _, files in os.walk(cache_dir):
+        for file in files:
+            if file.startswith("temp_"):
+                file_path = os.path.join(root, file)
+                try:
+                    file_mod_time = os.path.getmtime(file_path)
+                    if time.time() - file_mod_time > temp_threshold:
+                        os.remove(file_path)
+                        files_removed += 1
+                        logger.info(f"Removed temporary file: {file_path}")
+                except Exception as e:
+                    logger.error(f"Failed to remove temporary file {file_path}: {str(e)}")
 
     return updated_media_files, files_removed
 
@@ -266,7 +277,11 @@ async def download_new_files(media_files: list, cache_dir: str):
                 logger.error(f"Invalid file data: {file_data}")
                 continue
                 
-            cache_path = os.path.join(cache_dir, str(channel), str(post_id), file_unique_id)
+            channel_dir = os.path.join(cache_dir, str(channel))
+            post_dir = os.path.join(channel_dir, str(post_id))
+            os.makedirs(post_dir, exist_ok=True)
+            
+            cache_path = os.path.join(post_dir, file_unique_id)
             
             if not os.path.exists(cache_path):
                 files_to_download += 1
diff --git a/post_parser.py b/post_parser.py
index c89b74e..fa8a55a 100644
--- a/post_parser.py
+++ b/post_parser.py
@@ -59,12 +59,26 @@ class PostParser:
         self.allowed_tags = ['p', 'a', 'b', 'i', 'strong', 'em', 'ul', 'ol', 'li', 'br', 'div', 'span', 'img', 'video', 'audio', 'source']
         self.allowed_attributes = {
             'a': ['href', 'title', 'target'],
-            'img': ['src', 'alt', 'style'],
-            'video': ['controls', 'src', 'style'],
-            'audio': ['controls', 'style'],
+            'img': {
+                'src': True,
+                'alt': True,
+                'style': ['max-width', 'max-height', 'object-fit']
+            },
+            'video': {
+                'controls': True,
+                'src': True,
+                'style': ['max-width', 'max-height']
+            },
+            'audio': {
+                'controls': True,
+                'style': ['width', 'max-width']
+            },
             'source': ['src', 'type'],
-            'div': ['class', 'style'],
-            'span': ['class'],
+            'div': {
+                'class': True,
+                'style': ['margin']
+            },
+            'span': ['class']
         }
 
     def _debug_message(self, message: Message) -> Message:
@@ -227,9 +241,6 @@ class PostParser:
                 tags=self.allowed_tags,
                 attributes=self.allowed_attributes,
                 strip=True,
-                css_sanitizer=bleach.css_sanitizer.CSSSanitizer(
-                    allowed_css_properties=['max-width', 'max-height', 'width', 'height', 'object-fit']
-                )
             )
         except Exception as e:
             logger.error(f"html_sanitization_error: {str(e)}")