Update app.py

app.py

@@ -1,7 +1,7 @@
 from flask import Flask, send_file, request, make_response, Response
 import yt_dlp
 import os
-from urllib.parse import urlparse
+from urllib.parse import urlparse, unquote
 from gallery_dl import job
 import requests
 from datetime import datetime
@@ -9,6 +9,7 @@ import tempfile
 import shutil
 import mimetypes
 import uuid
+import re
 
 app = Flask(__name__)
 
@@ -24,11 +25,29 @@ class MediaDownloader:
         os.makedirs(self.video_path, exist_ok=True)
 
     def cleanup(self):
-        shutil.rmtree(self.temp_dir)
+        try:
+            if os.path.exists(self.temp_dir):
+                shutil.rmtree(self.temp_dir)
+        except Exception as e:
+            print(f"Cleanup error: {str(e)}")
+
+    def is_valid_url(self, url):
+        try:
+            result = urlparse(url)
+            return all([result.scheme, result.netloc])
+        except:
+            return False
+
+    def sanitize_filename(self, filename):
+        return re.sub(r'[<>:"/\\|?*]', '_', filename)
 
     def download_video(self, url):
+        if not self.is_valid_url(url):
+            return None
+
         unique_id = str(uuid.uuid4())
         output_template = os.path.join(self.video_path, f'{unique_id}.%(ext)s')
+        
         ydl_opts = {
             'format': 'best',
             'outtmpl': output_template,
@@ -47,11 +66,15 @@ class MediaDownloader:
                 if info:
                     filename = ydl.prepare_filename(info)
                     return filename if os.path.exists(filename) else None
+        except Exception as e:
+            print(f"Video download error: {str(e)}")
             return None
-        except:
-            return None
+        return None
 
     def download_images(self, url):
+        if not self.is_valid_url(url):
+            return None
+
         try:
             class UrlDl(job.Job):
                 def __init__(self, url, parent=None):
@@ -72,14 +95,15 @@ class MediaDownloader:
                 try:
                     response = requests.get(img_url, headers={
                         'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
-                    })
+                    }, timeout=30)
                     response.raise_for_status()
 
                     unique_id = str(uuid.uuid4())
                     ext = os.path.splitext(urlparse(img_url).path)[1]
-                    if not ext:
+                    if not ext or ext.lower() not in ['.jpg', '.jpeg', '.png', '.gif', '.webp']:
                         ext = '.jpg'
-                    filename = f"{unique_id}{ext}"
+                    
+                    filename = self.sanitize_filename(f"{unique_id}{ext}")
                     filepath = os.path.join(self.image_path, filename)
 
                     with open(filepath, 'wb') as f:
@@ -97,6 +121,9 @@ class MediaDownloader:
             return None
 
     def download_media(self, url):
+        if not url or not self.is_valid_url(url):
+            return None
+
         video_path = self.download_video(url)
         if video_path:
             return [video_path]
@@ -114,8 +141,8 @@ def health_check():
 @app.route('/')
 def home():
     return """
-    <h1>downloader</h1>
-    <p>Use: /download?url=url</p>
+    <h1>Downloader</h1>
+    <p>Use: /download?url=encoded_url</p>
     """
 
 @app.route('/download')
@@ -125,45 +152,55 @@ def download():
         if not url:
             return "No URL", 400
 
+        url = unquote(url)
         downloader = MediaDownloader()
 
         try:
             files = downloader.download_media(url)
 
-            if files and len(files) > 0:
-                if len(files) == 1:
-                    response = make_response(send_file(files[0], as_attachment=True))
-                    @response.call_on_close
-                    def cleanup():
-                        downloader.cleanup()
-                    return response
-                else:
-                    def generate():
+            if not files:
+                downloader.cleanup()
+                return "Download failed", 400
+
+            if len(files) == 1:
+                response = make_response(send_file(
+                    files[0],
+                    as_attachment=True,
+                    download_name=os.path.basename(files[0])
+                ))
+                @response.call_on_close
+                def cleanup():
+                    downloader.cleanup()
+                return response
+            else:
+                def generate():
+                    try:
                         for file_path in files:
-                            with open(file_path, 'rb') as f:
-                                content = f.read()
-                                filename = os.path.basename(file_path)
-                                yield f'Content-Disposition: attachment; filename="{filename}"\n'.encode()
-                                yield f'Content-Type: {mimetypes.guess_type(file_path)[0]}\n'.encode()
-                                yield f'Content-Length: {len(content)}\n\n'.encode()
-                                yield content
-                                yield b'\n--boundary--\n'
+                            if os.path.exists(file_path):
+                                with open(file_path, 'rb') as f:
+                                    content = f.read()
+                                    filename = os.path.basename(file_path)
+                                    yield f'--boundary\n'.encode()
+                                    yield f'Content-Disposition: attachment; filename="{filename}"\n'.encode()
+                                    yield f'Content-Type: {mimetypes.guess_type(file_path)[0]}\n'.encode()
+                                    yield f'Content-Length: {len(content)}\n\n'.encode()
+                                    yield content
+                                    yield b'\n'
+                        yield b'--boundary--\n'
+                    finally:
                         downloader.cleanup()
 
-                    response = Response(
-                        generate(),
-                        mimetype='multipart/x-mixed-replace; boundary=boundary',
-                        direct_passthrough=True
-                    )
-                    response.headers['Content-Type'] = 'multipart/x-mixed-replace; boundary=boundary'
-                    return response
-            else:
-                downloader.cleanup()
-                return "failed", 400
+                response = Response(
+                    generate(),
+                    mimetype='multipart/x-mixed-replace; boundary=boundary',
+                    direct_passthrough=True
+                )
+                response.headers['Content-Type'] = 'multipart/x-mixed-replace; boundary=boundary'
+                return response
 
         except Exception as e:
             downloader.cleanup()
-            return f"Error during download: {str(e)}", 500
+            return f"Download error: {str(e)}", 500
 
     except Exception as e:
         return f"Server error: {str(e)}", 500