Update main.py

main.py

@@ -15,12 +15,17 @@ from telethon.errors import FloodWaitError
 from telethon.tl.types import DocumentAttributeVideo
 
 # 🔹 Configuration
+# Best practice: Use environment variables, with fallbacks for easy testing
 API_ID = int(os.environ.get("API_ID", 20687211))
 API_HASH = os.environ.get("API_HASH", "4523f58b045175baaeaf1ba29733f31c")
 BOT_TOKEN = os.environ.get("BOT_TOKEN", "8318388017:AAGCgx_XJqBcyOpdt-5DlMNTS7wU9Qk6J4c")
+# Set this env variable to your own Telegram User ID
 BOT_OWNER_ID = int(os.environ.get("BOT_OWNER_ID", 7014665654))
 
-DOWNLOAD_DIR = "downloads" # Safe download location
+# --- File Management ---
+# We use a /data directory which is created in the Dockerfile for persistent storage
+SESSION_FILE = "/data/MegaNZBot"
+DOWNLOAD_DIR = "downloads" # This will be created inside /app
 
 # 🔹 Configure Logging
 logging.basicConfig(
@@ -33,13 +38,14 @@ logging.getLogger('telethon').setLevel(logging.WARNING)
 
 # 🔹 Initialize
 try:
-    bot = TelegramClient("MegaNZBot", API_ID, API_HASH)
+    # Use the /data path for the session file to fix permission errors in Docker
+    bot = TelegramClient(SESSION_FILE, API_ID, API_HASH)
     mega = Megatools()
 except Exception as e:
     logger.critical(f"Failed to initialize bot or Megatools: {e}")
     exit(1)
 
-# --- Core Helper Functions ---
+# --- Helper Functions ---
 
 def progress_bar(percentage, bar_length=20):
     """Generates a text-based progress bar."""
@@ -84,53 +90,52 @@ def clean_directory(directory):
 
 async def get_video_duration(file_path):
     """Tries to get video duration using ffprobe."""
-    command = [
-        "ffprobe", "-v", "quiet", "-print_format", "json",
-        "-show_format", file_path
-    ]
     try:
+        command = [
+            "ffprobe", "-v", "quiet", "-print_format", "json",
+            "-show_format", file_path
+        ]
         process = await asyncio.create_subprocess_exec(
             *command,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE
         )
         stdout, stderr = await process.communicate()
+        
         if process.returncode == 0:
             data = json.loads(stdout)
             return int(float(data['format']['duration']))
         else:
-            logger.warning(f"ffprobe failed: {stderr.decode()}")
+            logger.warning(f"ffprobe failed (duration): {stderr.decode()}")
             return 0
     except FileNotFoundError:
-        logger.warning("ffprobe not found. Cannot get video duration.")
+        logger.warning("ffprobe not found. Cannot get video duration. (Did you install ffmpeg?)")
         return 0
     except Exception as e:
-        logger.error(f"Error getting video duration: {e}")
+        logger.error(f"Error in get_video_duration: {e}")
         return 0
 
-async def upload_file(chat, file_path, filename, status_msg, status_prefix=""):
+# --- End of Helper Functions ---
+
+
+async def upload_file(chat, status_msg, file_path, filename, file_index, total_files):
     """
-    Handles the entire upload process for a single file:
-    - Generates caption
-    - Shows upload progress
-    - Gets video duration
-    - Sends the file
-    - Deletes the local file
+    Handles the upload process for a single file, including progress,
+    captions, and cleanup.
     """
     try:
         size_bytes = os.path.getsize(file_path)
         size_mb = round(size_bytes / (1024 * 1024), 2)
         ext = os.path.splitext(filename)[1].lower() or ".file"
-
+        
         caption = (
             f"📂 **File:** `{filename}`\n"
             f"📦 **Size:** {size_mb} MB\n"
             f"📄 **Extension:** `{ext}`"
         )
-        
-        status_text_prefix = f"📤 **Uploading {status_prefix}:**\n`{filename}`"
-        await edit_message(status_msg, f"{status_text_prefix}\n{progress_bar(0)} 0%")
 
+        await edit_message(status_msg, f"📤 **Uploading file {file_index}/{total_files}:**\n`{filename}`\n{progress_bar(0)} 0%")
+        
         last_update_time = 0
         async def progress_callback(current, total):
             nonlocal last_update_time
@@ -139,144 +144,142 @@ async def upload_file(chat, file_path, filename, status_msg, status_prefix=""):
             last_update_time = current_time
             percent = int(current * 100 / total)
             bar = progress_bar(percent)
-            await edit_message(status_msg, f"{status_text_prefix}\n{bar} {percent}%")
+            await edit_message(status_msg, f"📤 **Uploading file {file_index}/{total_files}:**\n`{filename}`\n{bar} {percent}%")
 
+        # Handle file types
         is_video = ext in [".mp4", ".mov", ".mkv", ".avi"]
-        is_image = ext in [".jpg", ".jpeg", ".png", ".webp"]
         
         attributes = []
         if is_video:
             duration = await get_video_duration(file_path)
             attributes.append(DocumentAttributeVideo(duration=duration, w=0, h=0, supports_streaming=True))
-
+        
         await bot.send_file(
             chat,
             file_path,
             caption=caption,
             attributes=attributes or None,
-            spoiler=True, # Send all files with spoiler
-            force_document=not (is_video or is_image),
+            spoiler=True,
+            force_document=not is_video, # Send as document if not a video
             progress_callback=progress_callback
         )
-        logger.info(f"Uploaded {filename}.")
         
+        logger.info(f"Uploaded {filename}.")
+        return True
+    
     except Exception as e:
         logger.error(f"Failed to upload file {file_path}: {e}", exc_info=True)
-        await status_msg.reply(f"⚠️ Failed to upload file: `{filename}`. Error: {e}")
+        await bot.send_message(chat, f"⚠️ Failed to upload file: `{filename}`. Error: {e}")
+        return False
+    
     finally:
-        if os.path.exists(file_path): 
-            os.remove(file_path) # Delete file after upload
-
-# --- End of Helper Functions ---
-
+        # Clean up individual file after upload
+        if os.path.exists(file_path):
+            os.remove(file_path)
 
 @bot.on(events.NewMessage(pattern="(?i)https://mega\\.nz/"))
 async def handle_mega(event):
     message_text = event.message.text
     chat = event.chat_id
     user_id = event.sender_id
-
+    
     status_msg = await event.reply("🔍 **Processing your link...**")
-
+    
     match = re.search(r"(https://mega\.nz/[^\s)\]]+)", message_text)
+    
     if not match:
-        logger.warning(f"Could not parse a mega.nz link from: {message_text}")
         await edit_message(status_msg, "❌ **Error:** Could not find a valid mega.nz link.")
         await delete_message(status_msg, delay=10)
         return
-
+        
     url = match.group(1).strip()
     
-    # Use a unique, safe download directory for this job
+    # Sandboxed directory for this specific job
     job_id = str(uuid.uuid4())
     download_dir = os.path.join(DOWNLOAD_DIR, f"user_{user_id}", job_id)
     os.makedirs(download_dir, exist_ok=True)
     
     try:
-        if "/folder/" in url:
-            # --- FOLDER LOGIC (Download all, then upload one-by-one) ---
+        is_folder = "/folder/" in url
+
+        if is_folder:
+            # --- FOLDER LOGIC ---
             logger.info(f"User {user_id} sent folder link: {url}")
             await edit_message(status_msg, "📁 **Folder link detected.**\n⬇️ Downloading all files... (This may take a while)")
-
+            
             try:
-                # This is the blocking call that works in your environment
+                # This is a blocking call. It will download the entire folder.
+                # This is the only method that works in your environment.
                 mega.download(url, path=download_dir)
                 logger.info(f"Folder downloaded to: {download_dir}")
             except Exception as e:
                 logger.error(f"Folder download failed for {url}: {e}", exc_info=True)
-                await edit_message(status_msg, f"❌ **Error:** Folder download failed:\n`{e}`")
-                await delete_message(status_msg, delay=10)
-                clean_directory(download_dir)
-                return
+                raise Exception(f"Folder download failed: {e}")
 
-            await edit_message(status_msg, "✅ **Download complete!**\nPreparing to upload files...")
+            await edit_message(status_msg, "✅ **Download complete!**\nPreparing to upload...")
 
-            uploaded_count = 0
+            # --- Upload 1-by-1 ---
             files_to_upload = []
             for root, dirs, files in os.walk(download_dir):
                 for file in files:
                     files_to_upload.append((os.path.join(root, file), file))
-
+            
             if not files_to_upload:
                 await edit_message(status_msg, "ℹ️ The folder appears to be empty.")
                 await delete_message(status_msg, delay=10)
-                clean_directory(download_dir)
                 return
 
             await edit_message(status_msg, f"Found {len(files_to_upload)} files. Starting upload...")
-
-            for i, (file_path, filename) in enumerate(files_to_upload):
-                # Use the helper function to upload
-                await upload_file(
-                    chat,
-                    file_path,
-                    filename,
-                    status_msg,
-                    status_prefix=f"file {i+1}/{len(files_to_upload)}"
-                )
-                uploaded_count += 1
-                # The helper function now handles file deletion
-
+            
+            uploaded_count = 0
+            for i, (file_path, file_name) in enumerate(files_to_upload):
+                if await upload_file(chat, status_msg, file_path, file_name, i+1, len(files_to_upload)):
+                    uploaded_count += 1
+            
             await edit_message(status_msg, f"✅ **Upload complete!**\nSuccessfully sent {uploaded_count} file(s).")
-            await delete_message(status_msg, delay=10)
 
         else:
             # --- SINGLE FILE LOGIC ---
             logger.info(f"User {user_id} sent file link: {url}")
             await edit_message(status_msg, "⏳ **Fetching file info...**")
+            
+            try:
+                filename = mega.filename(url)
+            except MegaError as e:
+                logger.error(f"MegaError (filename): {e}", exc_info=True)
+                await edit_message(status_msg, f"❌ **Error:** Link seems invalid or expired.\n`{e}`")
+                return
 
-            filename = mega.filename(url)
             if not filename:
-                logger.warning(f"Could not fetch filename for: {url}")
-                await edit_message(status_msg, "❌ **Error:** Could not fetch filename. The link might be invalid.")
-                await delete_message(status_msg, delay=10)
+                await edit_message(status_msg, "❌ **Error:** Could not fetch filename.")
                 return
 
-            temp_path = os.path.join(download_dir, filename)
+            file_path = os.path.join(download_dir, filename)
             await edit_message(status_msg, f"⬇️ **Downloading:**\n`{filename}`")
 
             try:
-                # Use download for single files as it's reliable
-                mega.download(url, path=temp_path)
-                logger.info(f"Downloaded to: {temp_path}")
+                # Blocking call to download the single file
+                mega.download(url, path=file_path)
+                logger.info(f"File downloaded to: {file_path}")
             except Exception as e:
-                logger.error(f"Download failed for {url}: {e}", exc_info=True)
-                await edit_message(status_msg, f"❌ **Error:** Download failed:\n`{e}`")
-                await delete_message(status_msg, delay=10)
-                return # temp_path will be cleaned in finally
+                logger.error(f"File download failed for {url}: {e}", exc_info=True)
+                raise Exception(f"File download failed: {e}")
 
-            # Use the upload helper function
-            await upload_file(chat, temp_path, filename, status_msg)
-            await delete_message(status_msg, delay=0) # Delete status msg immediately
+            await upload_file(chat, status_msg, file_path, filename, 1, 1)
+            await edit_message(status_msg, "✅ **Upload complete!**")
+        
+        await delete_message(status_msg, delay=10)
 
     except Exception as e:
-        logger.error(f"Error processing file {url}: {e}", exc_info=True)
+        logger.error(f"An unexpected error occurred processing {url}: {e}", exc_info=True)
         await edit_message(status_msg, f"⚠️ **An unexpected error occurred:**\n`{e}`")
         await delete_message(status_msg, delay=10)
+    
     finally:
         # Clean up the entire job directory
         clean_directory(download_dir)
 
+# --- Standard Bot Commands ---
 
 @bot.on(events.NewMessage(pattern="/start"))
 async def start(event):
@@ -304,11 +307,10 @@ async def help_command(event):
         "**How to use:**\n"
         "1.  Send me any public `mega.nz` file link.\n"
         "2.  Send me any public `mega.nz` folder link.\n"
-        "3.  I will download the file(s) and upload them to this chat.\n\n"
+        "3.  I will download the file(s) and upload them directly to this chat.\n\n"
         "**Features:**\n"
         "✅  Upload progress.\n"
         "✅  Supports both files and folders.\n"
-        "✅  Folder files are sent 1-by-1.\n"
         "✅  Gets video duration (if `ffprobe` is installed).\n"
         "✅  Cleans up files after upload to save space."
     )