Create telegram_bot.py

telegram_bot.py

@@ -0,0 +1,581 @@
+import os
+import logging
+from pyrogram import Client, filters, enums
+import requests
+import tempfile
+from urllib.parse import urljoin
+from config import API_ID, API_HASH, BOT_TOKEN, API_BASE_URL
+import time
+import math
+import asyncio
+import re
+from collections import defaultdict
+from datetime import datetime, timedelta
+import urllib3
+from asyncio import Lock
+
+# Enable logging
+logging.basicConfig(
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    level=logging.INFO
+)
+logger = logging.getLogger(__name__)
+
+# Disable SSL warnings
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+# Initialize the Pyrogram Client first
+app = Client(
+    "file_sharing_bot",
+    api_id=API_ID,
+    api_hash=API_HASH,
+    bot_token=BOT_TOKEN,
+    in_memory=True
+)
+
+# Set timeout for requests
+TIMEOUT = 30  # seconds
+
+# User-specific progress tracking
+user_progress = defaultdict(dict)
+
+# Rate limiting per user
+rate_limit = defaultdict(lambda: {'last_update': 0, 'count': 0})
+
+# Add after other global variables
+upload_locks = {}
+
+class RateLimiter:
+    def __init__(self, interval=1):
+        self.interval = interval
+        self.last_check = defaultdict(float)
+        
+    async def can_proceed(self, user_id):
+        now = time.time()
+        if now - self.last_check[user_id] < self.interval:
+            return False
+        self.last_check[user_id] = now
+        return True
+
+rate_limiter = RateLimiter(interval=2)
+
+async def progress(current, total, message, start_time, action="Uploading", user_id=None):
+    """Progress callback with per-user rate limiting"""
+    try:
+        if not user_id:
+            user_id = message.chat.id
+            
+        now = time.time()
+        
+        # Check rate limit for this user
+        if not await rate_limiter.can_proceed(user_id):
+            return
+            
+        elapsed_time = now - start_time
+        if elapsed_time == 0:
+            return
+
+        # Store progress for this user
+        user_progress[user_id] = {
+            'current': current,
+            'total': total,
+            'speed': current / elapsed_time,
+            'progress': (current * 100) / total
+        }
+        
+        progress_data = user_progress[user_id]
+        
+        # Format progress bar
+        bar_length = 20
+        filled_length = int(progress_data['progress'] / 100 * bar_length)
+        bar = '█' * filled_length + '░' * (bar_length - filled_length)
+        
+        # Calculate ETA
+        eta_seconds = (total - current) / progress_data['speed'] if progress_data['speed'] > 0 else 0
+        
+        text = (
+            f"{action} File...\n"
+            f"[{bar}] {progress_data['progress']:.1f}%\n"
+            f"Size: {format_size(current)}/{format_size(total)}\n"
+            f"Speed: {format_size(progress_data['speed'])}/s\n"
+            f"ETA: {int(eta_seconds)}s"
+        )
+        
+        try:
+            await message.edit_text(text)
+        except Exception as e:
+            logger.debug(f"Progress update failed: {e}")
+            
+    except Exception as e:
+        logger.error(f"Progress error for user {user_id}: {e}")
+
+# Add this class after other imports
+class ProgressTracker:
+    def __init__(self, message, action="Uploading"):
+        self.message = message
+        self.action = action
+        self.start_time = time.time()
+        self.last_update_time = 0
+        self.edit_failed = False
+        
+    async def update(self, current, total):
+        now = time.time()
+        # Update only every 1 second
+        if now - self.last_update_time < 1:
+            return
+            
+        self.last_update_time = now
+        elapsed_time = now - self.start_time
+        if elapsed_time == 0:
+            return
+            
+        speed = current / elapsed_time
+        progress_percent = (current * 100) / total
+        
+        # Calculate ETA
+        remaining_bytes = total - current
+        eta_seconds = remaining_bytes / speed if speed > 0 else 0
+        
+        # Format progress bar
+        bar_length = 20
+        filled_length = int(progress_percent / 100 * bar_length)
+        bar = '█' * filled_length + '░' * (bar_length - filled_length)
+        
+        text = (
+            f"{self.action} to server...\n"
+            f"[{bar}] {progress_percent:.1f}%\n"
+            f"Size: {format_size(current)}/{format_size(total)}\n"
+            f"Speed: {format_size(speed)}/s\n"
+            f"ETA: {int(eta_seconds)}s"
+        )
+        
+        # If upload is complete, update the message
+        if current >= total:
+            text = "✅ Upload complete! Processing..."
+            
+        try:
+            if not self.edit_failed:
+                try:
+                    await self.message.edit_text(text)
+                except Exception as e:
+                    logger.debug(f"Edit failed, switching to new messages: {e}")
+                    self.edit_failed = True
+                    # Delete old message
+                    try:
+                        await self.message.delete()
+                    except:
+                        pass
+                    # Send new message
+                    self.message = await self.message.reply_text(text)
+            else:
+                # Delete old message and send new one
+                try:
+                    await self.message.delete()
+                except:
+                    pass
+                self.message = await self.message.reply_text(text)
+                
+        except Exception as e:
+            logger.debug(f"Progress update failed: {e}")
+
+class ProgressFile:
+    def __init__(self, file, size, progress_callback):
+        self.file = file
+        self.size = size
+        self.progress_callback = progress_callback
+        self.uploaded = 0
+        
+    def read(self, chunk_size=-1):
+        data = self.file.read(chunk_size)
+        if data:
+            self.uploaded += len(data)
+            asyncio.create_task(
+                self.progress_callback(self.uploaded, self.size)
+            )
+        return data
+        
+    def seek(self, offset, whence=0):
+        return self.file.seek(offset, whence)
+        
+    def tell(self):
+        return self.file.tell()
+        
+    def close(self):
+        return self.file.close()
+        
+    def fileno(self):
+        return self.file.fileno()
+        
+    def readable(self):
+        return True
+        
+    def seekable(self):
+        return True
+        
+    def writable(self):
+        return False
+
+# Modify document handler for concurrent processing
+@app.on_message(filters.document)
+async def handle_document(client, message):
+    user_id = message.from_user.id
+    
+    # Check if user is already uploading
+    if user_id in upload_locks:
+        await message.reply_text("⚠️ Please wait for your current upload to finish.")
+        return
+        
+    upload_locks[user_id] = Lock()
+    
+    try:
+        async with upload_locks[user_id]:
+            status_msg = await message.reply_text("Starting file processing...")
+            
+            # Create user-specific temp directory
+            user_temp_dir = os.path.join("temp", str(user_id))
+            os.makedirs(user_temp_dir, exist_ok=True)
+            
+            # Get original filename
+            original_filename = message.document.file_name
+            safe_filename = sanitize_filename(original_filename)
+            file_path = os.path.join(user_temp_dir, safe_filename)
+            
+            try:
+                # Download with user-specific progress tracking
+                start_time = time.time()
+                await message.download(
+                    file_name=file_path,
+                    progress=progress,
+                    progress_args=(status_msg, start_time, "Downloading", user_id)
+                )
+                
+                # Get file size for upload progress
+                file_size = os.path.getsize(file_path)
+                progress_tracker = ProgressTracker(status_msg, "Uploading")
+                
+                with open(file_path, 'rb') as f:
+                    # Create a wrapper for the file to track upload progress
+                    progress_file = ProgressFile(f, file_size, progress_tracker.update)
+                    files = {'file': (original_filename, progress_file)}
+                    
+                    try:
+                        response = requests.post(
+                            f"{API_BASE_URL}/upload/",
+                            files=files,
+                            params={"user_id": str(user_id)},
+                            timeout=60,
+                            verify=False
+                        )
+                    
+                        if response.status_code == 200:
+                            result = response.json()
+                            file_url = urljoin(API_BASE_URL, result['access_code'])
+                            
+                            # Send a new message instead of editing
+                            await status_msg.delete()
+                            await message.reply_text(
+                                f"✅ File uploaded successfully!\n\n"
+                                f"📄 Filename: {original_filename}\n"
+                                f"🔑 Access Code: <code>{result['access_code']}</code>\n"
+                                f"🔗 Direct Link: {file_url}\n\n"
+                                f"Anyone with this link can download the file.",
+                                parse_mode=enums.ParseMode.HTML
+                            )
+                        else:
+                            error_msg = response.json().get('detail', 'Unknown error')
+                            await message.reply_text(f"❌ Upload failed: {error_msg}")
+                            
+                    except Exception as upload_error:
+                        logger.error(f"Upload error detail: {upload_error}")
+                        await message.reply_text("❌ Upload failed at final stage. Please try again.")
+                        
+            except Exception as e:
+                logger.error(f"File handling error: {e}")
+                await message.reply_text("❌ Error processing file. Please try again.")
+                
+    except Exception as e:
+        logger.error(f"Error for user {user_id}: {e}")
+        await message.reply_text("❌ Sorry, something went wrong. Please try again.")
+    finally:
+        # Cleanup user-specific temp files
+        try:
+            if os.path.exists(user_temp_dir):
+                for file in os.listdir(user_temp_dir):
+                    os.remove(os.path.join(user_temp_dir, file))
+                os.rmdir(user_temp_dir)
+        except Exception as e:
+            logger.error(f"Cleanup error for user {user_id}: {e}")
+            
+        # Remove the lock
+        if user_id in upload_locks:
+            del upload_locks[user_id]
+
+def format_size(size):
+    """Format size in bytes to human readable format"""
+    units = ['B', 'KB', 'MB', 'GB']
+    size = float(size)
+    unit = 0
+    while size >= 1024 and unit < len(units) - 1:
+        size /= 1024
+        unit += 1
+    return f"{size:.2f} {units[unit]}"
+
+@app.on_message(filters.command("start"))
+async def start_command(client, message):
+    """Handle the /start command"""
+    await message.reply_text(
+        'Hi! I can help you share files.\n'
+        'Just send me any file and I will give you a link to share it.\n\n'
+        'Commands:\n'
+        '/start - Show this help message\n'
+        '/list - List all uploaded files\n'
+        '/delete <code> - Delete a file using access code\n'
+        '/stats - View your usage statistics\n\n'
+        '💡 You can also send me an access code directly to get the file!'
+    )
+
+@app.on_message(filters.command("list"))
+async def list_command(client, message):
+    """Handle the /list command"""
+    try:
+        # Get user's Telegram ID
+        user_id = str(message.from_user.id)
+        
+        # Get only this user's files
+        response = requests.get(f"{API_BASE_URL}/files/{user_id}")
+        files = response.json()
+        
+        if not files['files']:
+            await message.reply_text("📂 You haven't uploaded any files yet.")
+            return
+
+        async def send_long_message(text, parse_mode=None):
+            MAX_LENGTH = 4000
+            messages = []
+            current_msg = "📂 Your Files:\n\n"
+            
+            for line in text.split('\n'):
+                if len(current_msg + line + '\n') > MAX_LENGTH:
+                    messages.append(current_msg)
+                    current_msg = "📂 Your Files (continued):\n\n" + line + '\n'
+                else:
+                    current_msg += line + '\n'
+            
+            if current_msg:
+                messages.append(current_msg)
+            
+            for i, msg_text in enumerate(messages, 1):
+                if len(messages) > 1:
+                    msg_text += f"\n📃 Page {i}/{len(messages)}"
+                await message.reply_text(msg_text, parse_mode=parse_mode)
+
+        # Prepare combined message with better formatting
+        files_msg = ""
+        for i, file in enumerate(files['files'], 1):
+            file_url = urljoin(API_BASE_URL, file['access_code'])
+            files_msg += f"<b>{i}. {file['filename']}</b>\n"
+            files_msg += f"   ├─ 🔗 <a href='{file_url}'>Direct Link</a>\n"
+            files_msg += f"   └─ 🔑 Code: <code>{file['access_code']}</code>\n\n"
+
+        # Send message with HTML formatting
+        await send_long_message(files_msg, parse_mode=enums.ParseMode.HTML)
+
+    except Exception as e:
+        logger.error(f"Error listing files: {e}")
+        await message.reply_text("❌ Sorry, couldn't fetch your files.")
+
+def sanitize_filename(filename):
+    """Remove invalid characters from filename"""
+    # Remove invalid characters
+    filename = re.sub(r'[<>:"/\\|?*]', '', filename)
+    # Remove any leading/trailing spaces and dots
+    filename = filename.strip('. ')
+    # If filename is empty after sanitization, use a default name
+    if not filename:
+        filename = 'downloaded_file'
+    return filename
+
+@app.on_message(filters.command("delete"))
+async def delete_command(client, message):
+    """Handle the /delete command"""
+    try:
+        # Check if access code is provided
+        command_parts = message.text.split()
+        if len(command_parts) != 2:
+            await message.reply_text(
+                "❌ Please provide an access code.\n"
+                "Usage: /delete <access_code>"
+            )
+            return
+            
+        user_id = str(message.from_user.id)
+        access_code = command_parts[1]
+        
+        # Try to delete the file
+        response = requests.delete(
+            f"{API_BASE_URL}/delete/{access_code}",
+            params={"user_id": user_id},  # Add user_id to verify ownership
+            timeout=TIMEOUT,
+            verify=False
+        )
+        
+        if response.status_code == 200:
+            await message.reply_text("✅ File deleted successfully!")
+        else:
+            error_msg = response.json().get('detail', 'Unknown error')
+            await message.reply_text(f"❌ Error: {error_msg}")
+            
+    except Exception as e:
+        logger.error(f"Error deleting file: {e}")
+        await message.reply_text("❌ Sorry, couldn't delete the file.")
+
+@app.on_message(filters.command(["stats", "stats@your_bot_username"]))  # Add your bot's username
+async def stats_command(client, message):
+    """Handle the /stats command"""
+    print("stats command received")  # Debug print
+    
+    # Check if it's a private chat
+    if message.chat.type != enums.ChatType.PRIVATE:
+        await message.reply_text("Please use this command in private chat.")
+        return
+        
+    try:
+        user_id = str(message.from_user.id)
+        print(f"Processing stats for user_id: {user_id}")  # Debug print
+        
+        # Get stats from server
+        response = requests.get(
+            f"{API_BASE_URL}/stats/{user_id}", 
+            timeout=TIMEOUT,
+            verify=False
+        )
+        print(f"Server response: {response.text}")  # Debug print
+        
+        if response.status_code == 200:
+            stats = response.json()
+            
+            # Format the statistics message with better error handling
+            uploads = stats.get('uploads', 0)
+            downloads = stats.get('downloads', 0)
+            bytes_uploaded = stats.get('bytes_uploaded', 0)
+            bytes_downloaded = stats.get('bytes_downloaded', 0)
+            last_activity = stats.get('last_activity', 'No activity')
+            
+            stats_msg = (
+                "📊 Your File Sharing Statistics\n\n"
+                f"📤 Uploads: {uploads}\n"
+                f"📥 Downloads: {downloads}\n"
+                f"📈 Total Uploaded: {format_size(bytes_uploaded)}\n"
+                f"📉 Total Downloaded: {format_size(bytes_downloaded)}\n"
+                f"🕒 Last Activity: {last_activity}"
+            )
+            
+            await message.reply_text(stats_msg)
+        else:
+            logger.error(f"Stats error: {response.text}")  # Add error logging
+            await message.reply_text("❌ Couldn't fetch your statistics.")
+            
+    except Exception as e:
+        logger.error(f"Error fetching stats: {e}")
+        await message.reply_text("❌ Sorry, something went wrong while fetching your statistics.")
+
+@app.on_message(filters.text & filters.private & ~filters.via_bot & ~filters.forwarded)
+async def handle_text(client, message):
+    """Handle text messages as potential access codes"""
+    if message.text.startswith('/'):
+        return
+        
+    try:
+        user_id = str(message.from_user.id)
+        access_code = message.text.strip()
+        if len(access_code) != 8:
+            return
+
+        status_msg = await message.reply_text("🔍 Fetching file...")
+        
+        response = requests.get(
+            f"{API_BASE_URL}/download/{access_code}", 
+            stream=True, 
+            timeout=TIMEOUT
+        )
+        
+        if response.status_code == 200:
+            # Get and sanitize filename from headers
+            content_disposition = response.headers.get('content-disposition', '')
+            if 'filename=' in content_disposition:
+                filename = content_disposition.split('filename=')[-1].strip('"\'')
+                # URL decode the filename
+                filename = requests.utils.unquote(filename)
+            else:
+                filename = f"file_{access_code}"
+            
+            # Sanitize the filename
+            filename = sanitize_filename(filename)
+            
+            # Save file temporarily
+            temp_path = os.path.join("temp", filename)
+            os.makedirs("temp", exist_ok=True)
+            
+            try:
+                #Download file with progress
+                total_size = int(response.headers.get('content-length', 0))
+                current_size = 0
+                start_time = time.time()
+                
+                with open(temp_path, 'wb') as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        if chunk:
+                            f.write(chunk)
+                            current_size += len(chunk)
+                            await progress(
+                                current_size,
+                                total_size,
+                                status_msg,
+                                start_time,
+                                "Downloading"
+                            )
+                
+                # Send file to user
+                await status_msg.edit_text("📤 Sending file to you...")
+                await message.reply_document(
+                    temp_path,
+                    caption=f"📄 File fetched using access code: {access_code}"
+                )
+                await status_msg.delete()
+                
+                # Log the download
+                file_size = int(response.headers.get('content-length', 0))
+                requests.post(
+                    f"{API_BASE_URL}/log_download",
+                    json={
+                        "user_id": user_id,
+                        "file_size": file_size,
+                        "filename": filename
+                    },
+                    timeout=TIMEOUT
+                )
+                
+            finally:
+                # Clean up
+                try:
+                    os.remove(temp_path)
+                except:
+                    pass
+                    
+        else:
+            error_msg = response.json().get('detail', 'Unknown error')
+            await status_msg.edit_text(f"❌ Invalid access code or file not found")
+            
+    except Exception as e:
+        logger.error(f"Error fetching file: {e}")
+        await message.reply_text("❌ Sorry, something went wrong while fetching the file.")
+
+
+
+def main():
+    """Start the bot"""
+    print("Starting bot...")
+    app.run()
+
+if __name__ == '__main__':
+    main()