terabox_utils.py

# terabox_utils.py

import re
import requests
import asyncio
import os
import time
import math
import logging
from functools import partial
from typing import Optional, Tuple

import config

logger = logging.getLogger(__name__)
os.makedirs("downloads", exist_ok=True)

# --- Utility Functions ---
def format_bytes(size_bytes: int) -> str:
    if size_bytes <= 0:
        return "0 B"
    size_name = ("B", "KB", "MB", "GB", "TB")
    i = min(int(math.log(size_bytes, 1024)), len(size_name) - 1)
    p = math.pow(1024, i)
    s = round(size_bytes / p, 2)
    return f"{s} {size_name[i]}"

def format_eta(seconds_left: float) -> str:
    seconds_left = int(seconds_left)
    h, remainder = divmod(seconds_left, 3600)
    m, s = divmod(remainder, 60)
    if h > 0:
        return f"{h}h {m}m {s}s"
    elif m > 0:
        return f"{m}m {s}s"
    else:
        return f"{s}s"

# --- Short ID Extract ---
async def extract_terabox_short_id(full_url: str) -> Optional[str]:
    patterns = [
        r'terabox\.com/s/([a-zA-Z0-9_-]+)',
        r'teraboxapp\.com/s/([a-zA-Z0-9_-]+)',
        r'1024tera\.com/s/([a-zA-Z0-9_-]+)',
        r'freeterabox\.com/s/([a-zA-Z0-9_-]+)',
        r'terabox\.com/sharing/link\?surl=([a-zA-Z0-9_-]+)',
        r'terasharelink\.com/s/([a-zA-Z0-9_-]+)',
        r'4funbox\.com/s/([a-zA-Z0-9_-]+)',
        r'box-links\.com/s/([a-zA-Z0-9_-]+)'
    ]
    for p in patterns:
        if m := re.search(p, full_url, re.I):
            return m.group(1)
    return None

# --- Get Final URL and Filename ---
async def get_final_url_and_filename(original_link: str) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    payload = {"link": original_link}
    headers = {"User-Agent": "Mozilla/5.0"}

    try:
        loop = asyncio.get_event_loop()
        r = await loop.run_in_executor(
            None,
            partial(requests.post, config.TERABOX_WORKER_URL, headers=headers, json=payload, timeout=30)
        )
        r.raise_for_status()
        data = r.json()

        dl = data.get("proxy_url")
        fn = data.get("file_name")

        if data.get("error") or not dl or not fn:
            return None, None, data.get('error', 'Worker returned incomplete data.')

        return dl, fn, None

    except Exception as e:
        return None, None, str(e)

# --- Download with Progress ---
async def download_terabox_file(bot_instance, chat_id, msg_id, url, filename) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    safe_fn = re.sub(r'[\\/*?:"<>|]', "_", filename)[:200]
    download_path = os.path.join("downloads", f"{chat_id}_{time.time()}_{safe_fn}")

    try:
        loop = asyncio.get_event_loop()

        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
            "Accept": "*/*",
            "Referer": "https://teraboxapp.com/"
        }

        r = await loop.run_in_executor(
            None,
            partial(requests.get, url, headers=headers, stream=True, timeout=(10, 300), allow_redirects=True)
        )
        r.raise_for_status()

        total_size = int(r.headers.get('content-length', 0))
        dl_size = 0
        last_update = time.time()
        start_time = time.time()

        with open(download_path, 'wb') as f:
            for chunk in r.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    f.write(chunk)
                    dl_size += len(chunk)

                    elapsed = time.time() - start_time
                    speed = dl_size / elapsed if elapsed > 0 else 0
                    eta = (total_size - dl_size) / speed if speed > 0 else 0

                    pct = (dl_size / total_size * 100) if total_size > 0 else 0
                    prog_text = (
                        f"📥 <b>Downloading:</b> <code>{filename}</code>\n"
                        f"Progress: {format_bytes(dl_size)}/{format_bytes(total_size)} ({pct:.1f}%)\n"
                        f"Speed: {format_bytes(speed)}/s\n"
                        f"ETA: {format_eta(eta)}"
                    )
                    if time.time() - last_update > 2.5:
                        try:
                            await bot_instance.edit_message_text(
                                chat_id=chat_id,
                                message_id=msg_id,
                                text=prog_text,
                                parse_mode="HTML"
                            )
                        except Exception:
                            pass

                        last_update = time.time()

        # Optional thumbnail generation — skip if not a video
        video_exts = ('.mp4', '.mkv', '.mov', '.avi', '.webm')
        thumb_path = None
        if safe_fn.lower().endswith(video_exts):
            thumb_path = await generate_video_thumbnail(download_path, chat_id)

        return download_path, thumb_path, None

    except Exception as e:
        if os.path.exists(download_path):
            os.remove(download_path)
        return None, None, str(e)

# --- Generate Video Thumbnail (Optional) ---
async def generate_video_thumbnail(filepath: str, chat_id: int) -> Optional[str]:
    try:
        import subprocess
        thumb_path = os.path.join("downloads", f"{chat_id}_{time.time()}_thumb.jpg")

        ffmpeg_cmd = [
            "ffmpeg", "-y", "-i", filepath,
            "-ss", "00:00:01.000", "-vframes", "1", thumb_path
        ]
        subprocess.run(ffmpeg_cmd, check=True)

        if os.path.exists(thumb_path):
            return thumb_path
    except Exception as e:
        logger.warning(f"Thumbnail generation failed: {e}")

    return None