main.py

import requests
import random
import re
import json
import hashlib
import time
import psycopg2
from urllib.parse import urlparse, urljoin, urlencode, quote, unquote
from flask import Flask, request, jsonify, Response, stream_with_context

# Database configuration
DATABASE_URL = 'postgresql://neondb_owner:npg_yMhW5ZbeJPI6@ep-green-butterfly-adilhz1l-pooler.c-2.us-east-1.aws.neon.tech/neondb?sslmode=require&channel_binding=require'

class TeraBoxDownloader:
    def __init__(self, cookies=None):
        self.cache = {}
        self.cookies = cookies or []
        self.current_cookie_index = 0
        self.user_agent = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
        self.terabox_domains = [
            'terafileshare.com', 'www.terafileshare.com', 'terabox.com', 'www.terabox.com',
            'teraboxlink.com', '4funbox.com', 'www.4funbox.com', 'terasharelink.com',
            '1024terabox.com', 'www.1024terabox.com', 'terabox.app', 'www.terabox.app',
            'terabox.fun', '1024tera.com', '1024tera.co', '1024box.com', 'teraboxshare.com',
            'teraboxapp.com', 'momerybox.com', '4funbox.co', 'mirrobox.com', 'nephobox.com',
            'freeterabox.com', 'tibibox.com'
        ]
        self.init_db()
        self.load_cookies_from_db()

    def init_db(self):
        try:
            conn = psycopg2.connect(DATABASE_URL)
            cur = conn.cursor()
            cur.execute("""
                CREATE TABLE IF NOT EXISTS terabox_cookies (
                    id SERIAL PRIMARY KEY,
                    cookie_id TEXT UNIQUE NOT NULL,
                    host TEXT NOT NULL,
                    cookie TEXT NOT NULL
                )
            """)
            conn.commit()
            cur.close()
            conn.close()
        except Exception as e:
            print(f"DB Init Error: {e}")

    def load_cookies_from_db(self):
        try:
            conn = psycopg2.connect(DATABASE_URL)
            cur = conn.cursor()
            cur.execute("SELECT cookie_id, host, cookie FROM terabox_cookies")
            rows = cur.fetchall()
            if rows:
                self.cookies = [{'id': r[0], 'host': r[1], 'cookie': r[2]} for r in rows]
            cur.close()
            conn.close()
        except Exception as e:
            print(f"DB Load Error: {e}")

    def add_cookie_to_db(self, cookie_id, host, cookie_str):
        try:
            conn = psycopg2.connect(DATABASE_URL)
            cur = conn.cursor()
            # Explicitly update if exists to ensure the behavior requested
            cur.execute("""
                INSERT INTO terabox_cookies (cookie_id, host, cookie)
                VALUES (%s, %s, %s)
                ON CONFLICT (cookie_id) 
                DO UPDATE SET 
                    cookie = EXCLUDED.cookie,
                    host = EXCLUDED.host
            """, (str(cookie_id), host, cookie_str))
            conn.commit()
            cur.close()
            conn.close()
            self.load_cookies_from_db()
            return True
        except Exception as e:
            print(f"DB Add Error: {e}")
            return False

    def generate_hash_id(self, original_val):
        return hashlib.md5(f"{original_val}{time.time()}".encode()).hexdigest()

    def get_current_cookie_config(self):
        if not self.cookies:
            return None
        self.current_cookie_index = random.randint(0, len(self.cookies) - 1)
        return self.cookies[self.current_cookie_index]

    def make_request(self, url, selected_cookie_config, max_retries=3):
        if not selected_cookie_config:
            return {'ok': False, 'error': 'No cookies available'}
        last_error = None
        for attempt in range(max_retries):
            headers = {
                'User-Agent': self.user_agent,
                'Cookie': selected_cookie_config['cookie'],
                'Referer': url if 'terabox' in url else 'https://terabox.com/',
            }
            try:
                response = requests.get(url, headers=headers, timeout=15)
                if response.ok:
                    return {'body': response.text, 'status': response.status_code, 'ok': True, 'cookie_used': selected_cookie_config['id']}
                last_error = f"HTTP {response.status_code}"
            except Exception as e:
                last_error = str(e)
        return {'ok': False, 'error': last_error}

    def find_between(self, string, start, end):
        try:
            start_pos = string.index(start) + len(start)
            end_pos = string.index(end, start_pos)
            return string[start_pos:end_pos]
        except: return ""

    async def extract_js_token(self, cookie_config, surl):
        page_url = f"https://{cookie_config['host']}/sharing/link?surl={surl}"
        response = self.make_request(page_url, cookie_config)
        if not response['ok']: return None, "Failed"
        html = response['body']
        js_token = self.find_between(html, 'fn%28%22', '%22%29')
        if not js_token: js_token = self.find_between(html, 'fn("', '")')
        if not js_token: js_token = self.find_between(html, '"jsToken":"', '"')
        if not js_token:
            p = re.search(r'var\s+jsToken\s*=\s*[\'"]([^\'"]+)[\'"]', html)
            if p: js_token = p.group(1)
        return js_token, None

    async def get_streaming_info(self, cookie_config, uk, shareid, fid, sign, timestamp, js_token):
        """Fetch real m3u8 streaming URL from TeraBox"""
        try:
            streaming_api_url = f"https://{cookie_config['host']}/share/streaming?uk={uk}&shareid={shareid}&type=M3U8_AUTO_360&fid={fid}&sign={sign}&timestamp={timestamp}&jsToken={js_token}&esl=1&isplayer=1&ehps=1&clienttype=0&app_id=250528&web=1&channel=dubox"
            
            headers = {
                'User-Agent': self.user_agent,
                'Cookie': cookie_config['cookie'],
                'Referer': f"https://{cookie_config['host']}/sharing/link",
                'Origin': f"https://{cookie_config['host']}"
            }
            
            # This request returns the m3u8 content directly or a redirect
            response = requests.get(streaming_api_url, headers=headers, allow_redirects=False, timeout=10)
            
            if response.status_code in [301, 302]:
                return response.headers.get('Location'), None
            
            # If we get JSON, it might be an error or a direct stream URL in data
            try:
                data = response.json()
                if data.get('errno') == 2: # Invalid channel
                    alt_url = streaming_api_url.replace('clienttype=0', 'clienttype=1')
                    res2 = requests.get(alt_url, headers=headers, allow_redirects=False, timeout=10)
                    if res2.status_code in [301, 302]: return res2.headers.get('Location'), None
                    if res2.ok: return alt_url, None
                if data.get('errno'):
                    return None, f"API Error: {data.get('errno')}"
            except: pass

            if response.ok:
                return streaming_api_url, None
                
            return None, f"Streaming API failed: {response.status_code}"
        except Exception as e:
            return None, f"Streaming info exception: {str(e)}"

    async def process_url(self, url):
        """Process TeraBox URL and return JSON response"""
        if url in self.cache:
            cache_data, cache_time = self.cache[url]
            if time.time() - cache_time < 43200: # 12 hours
                return cache_data

        selected_cookie = self.get_current_cookie_config()
        if not selected_cookie: return {'error': 'No cookies configured'}
        
        surl = None
        surl_match = re.search(r'surl=([^&]+)', url)
        if surl_match:
            surl = surl_match.group(1)
        else:
            s_match = re.search(r'/s/([^/?&]+)', url)
            if s_match:
                surl = s_match.group(1)
                if surl.startswith('1') and len(surl) > 1:
                    surl = surl[1:]

        if not surl:
            return {'error': 'Invalid TeraBox URL - surl not found'}

        js_token, err = await self.extract_js_token(selected_cookie, surl)
        if not js_token:
            # Notify external API
            try: 
                # Send only the numeric part if it starts with 'cookie-'
                cid = selected_cookie['id']
                if isinstance(cid, str) and cid.startswith('cookie-'):
                    cid = cid.replace('cookie-', '')
                requests.get(f"https://b0da308e-dea2-4f7f-9716-e42f05617138-00-3v7r6qc4904uz.sisko.replit.dev:8000//regenerate?number={cid}", timeout=5)
            except: pass
            return {
                'error': 'Failed to extract jsToken',
                'detail': err,
                'cookie_used': selected_cookie['id']
            }

        params = {'clienttype': '5', 'jsToken': js_token, 'shorturl': surl, 'root': '1'}
        list_url = f"https://{selected_cookie['host']}/share/list?{urlencode(params)}"
        list_res = self.make_request(list_url, selected_cookie)

        if list_res['ok']:
            data = json.loads(list_res['body'])
            if data and data.get('errno') == 0:
                files = data.get('list', [])
                for f in files:
                    if 'thumbs' in f: del f['thumbs']

                if len(files) == 1 and str(files[0].get('isdir')) == '0':
                    f = files[0]
                    dlink = f.get('dlink')
                    if not dlink:
                        # Notify external API
                        try: 
                            cid = selected_cookie['id']
                            if isinstance(cid, str) and cid.startswith('cookie-'):
                                cid = cid.replace('cookie-', '')
                            requests.get(f"https://b0da308e-dea2-4f7f-9716-e42f05617138-00-3v7r6qc4904uz.sisko.replit.dev:8000//regenerate?number={cid}", timeout=5)
                        except: pass
                        return {'error': 'Download link not generated', 'cookie_used': selected_cookie['id']}

                    real_streaming_url, stream_err = await self.get_streaming_info(
                        selected_cookie, data.get('uk'), data.get('share_id'), 
                        f['fs_id'], data.get('sign'), data.get('server_time'), js_token
                    )
                    
                    hid = self.generate_hash_id(real_streaming_url or f['fs_id'])
                    
                    # Store stream info for proxy
                    self.cache[f"stream_{hid}"] = {
                        'url': real_streaming_url,
                        'cookie': selected_cookie['cookie']
                    }
                    # Store download info for proxy
                    self.cache[f"download_{hid}"] = {
                        'url': dlink,
                        'cookie': selected_cookie['cookie']
                    }
                    
                    result = {
                        'proxy_stream_url': f"/stream/{hid}.m3u8",
                        'proxy_download_url': f"/download/{hid}",
                        'cookie_used': selected_cookie['id']
                    }
                    self.cache[url] = (result, time.time())
                    return result
                    
                return {'error': 'Not a single file or directory'}
        
        # Notify external API
        try: 
            cid = selected_cookie['id']
            if isinstance(cid, str) and cid.startswith('cookie-'):
                cid = cid.replace('cookie-', '')
            requests.get(f"https://b0da308e-dea2-4f7f-9716-e42f05617138-00-3v7r6qc4904uz.sisko.replit.dev:8000//regenerate?number={cid}", timeout=5)
        except: pass
        return {'error': 'List failed', 'detail': list_res.get('error'), 'cookie_used': selected_cookie['id']}

    def proxy_download(self, hid):
        # We need to find the dlink for this hid
        # Since cache stores by URL, let's store dlink specifically
        download_data = self.cache.get(f"download_{hid}")
        if not download_data:
            return "Not Found", 404
            
        url = download_data['url']
        headers = {
            'User-Agent': self.user_agent,
            'Cookie': download_data['cookie'],
            'Referer': 'https://www.terabox.com/'
        }
        
        try:
            resp = requests.get(url, headers=headers, stream=True, timeout=30)
            excluded_headers = ['content-encoding', 'transfer-encoding', 'connection']
            headers = [(name, value) for (name, value) in resp.raw.headers.items()
                      if name.lower() not in excluded_headers]
            return Response(stream_with_context(resp.iter_content(chunk_size=32768)), 
                           resp.status_code, headers)
        except Exception as e:
            return str(e), 500

    def proxy_stream(self, hid):
        stream_data = self.cache.get(f"stream_{hid}")
        if not stream_data:
            return "Not Found", 404
            
        url = stream_data['url']
        headers = {
            'User-Agent': self.user_agent,
            'Cookie': stream_data['cookie'],
            'Referer': f"https://{urlparse(url).netloc}/",
            'Origin': f"https://{urlparse(url).netloc}"
        }
        
        try:
            # First, check if the URL is an m3u8 playlist or a redirect
            resp = requests.get(url, headers=headers, stream=True, timeout=15, allow_redirects=True)
            
            # If the response is an m3u8 playlist, we need to handle relative paths in it
            content_type = resp.headers.get('Content-Type', '').lower()
            is_m3u8 = 'mpegurl' in content_type or url.split('?')[0].endswith('.m3u8')
            
            if is_m3u8:
                playlist_content = resp.text
                # Use the final URL after redirects for the base URL
                base_url = resp.url.rsplit('/', 1)[0]
                
                # Replace relative segment URLs with absolute ones through our proxy
                def replace_segment(match):
                    segment_url = match.group(0).strip()
                    if not segment_url.startswith(('http://', 'https://')):
                        # Handle absolute paths starting with /
                        if segment_url.startswith('/'):
                            parsed = urlparse(resp.url)
                            full_segment_url = f"{parsed.scheme}://{parsed.netloc}{segment_url}"
                        else:
                            full_segment_url = f"{base_url}/{segment_url}"
                    else:
                        full_segment_url = segment_url
                    
                    # Store segment info in cache for proxying
                    segment_hid = hashlib.md5(full_segment_url.encode()).hexdigest()
                    self.cache[f"segment_{segment_hid}"] = {
                        'url': full_segment_url,
                        'cookie': stream_data['cookie']
                    }
                    return f"/segment/{segment_hid}.ts"

                # Match segments (usually any line that doesn't start with #)
                lines = playlist_content.splitlines()
                new_lines = []
                for line in lines:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        new_lines.append(replace_segment(re.match(r'.*', line)))
                    else:
                        new_lines.append(line)
                
                playlist_content = "\n".join(new_lines)
                r = Response(playlist_content, mimetype='application/vnd.apple.mpegurl')
                r.headers['Cache-Control'] = 'no-cache'
                return r

            excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
            headers = [(name, value) for (name, value) in resp.raw.headers.items()
                      if name.lower() not in excluded_headers]
            return Response(stream_with_context(resp.iter_content(chunk_size=8192)), 
                           resp.status_code, headers)
        except Exception as e:
            return str(e), 500

    def proxy_segment(self, hid):
        segment_data = self.cache.get(f"segment_{hid}")
        if not segment_data:
            return "Not Found", 404
            
        url = segment_data['url']
        headers = {
            'User-Agent': self.user_agent,
            'Cookie': segment_data['cookie'],
            'Referer': 'https://www.terabox.com/'
        }
        
        try:
            resp = requests.get(url, headers=headers, stream=True, timeout=30)
            excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
            headers = [(name, value) for (name, value) in resp.raw.headers.items()
                      if name.lower() not in excluded_headers]
            return Response(stream_with_context(resp.iter_content(chunk_size=16384)), 
                           resp.status_code, headers)
        except Exception as e:
            return str(e), 500

app = Flask(__name__)
dl = TeraBoxDownloader()

@app.route('/api/process')
async def process():
    url = request.args.get('url')
    if not url: return jsonify({'error': 'No URL'}), 400
    return jsonify(await dl.process_url(url))

@app.route('/stream/<hid>.m3u8')
def stream_proxy(hid):
    return dl.proxy_stream(hid)

@app.route('/download/<hid>')
def download_proxy(hid):
    return dl.proxy_download(hid)

@app.route('/segment/<hid>.ts')
def segment_proxy(hid):
    return dl.proxy_segment(hid)

@app.route('/api/cookies')
def add_cookie():
    # Delete cookie if delete parameter is provided
    delete_id = request.args.get('delete')
    if delete_id:
        try:
            conn = psycopg2.connect(DATABASE_URL)
            cur = conn.cursor()
            cur.execute("DELETE FROM terabox_cookies WHERE cookie_id = %s", (str(delete_id),))
            conn.commit()
            deleted = cur.rowcount
            cur.close()
            conn.close()
            dl.load_cookies_from_db()
            if deleted:
                return jsonify({'status': 'success', 'message': f'Cookie {delete_id} deleted'})
            return jsonify({'error': f'Cookie {delete_id} not found'}), 404
        except Exception as e:
            return jsonify({'error': str(e)}), 500

    # List cookies if no cookies parameter is provided
    if not request.args.get('cookies'):
        try:
            conn = psycopg2.connect(DATABASE_URL)
            cur = conn.cursor()
            cur.execute("SELECT cookie_id, host, cookie FROM terabox_cookies ORDER BY id")
            rows = cur.fetchall()
            cur.close()
            conn.close()
            
            cookie_list = []
            for r in rows:
                # Mask the middle part of the cookie for security while showing it's there
                c = r[2]
                masked = c[:10] + "..." + c[-10:] if len(c) > 20 else c
                cookie_list.append({
                    'number': r[0],
                    'host': r[1],
                    'cookie_preview': masked
                })
            return jsonify({
                'total_cookies': len(cookie_list),
                'cookies': cookie_list
            })
        except Exception as e:
            return jsonify({'error': str(e)}), 500

    cookie_str = request.args.get('cookies')
    number = request.args.get('number', f"cookie-{int(time.time())}")
    host = request.args.get('host', 'dm.terabox.app')
    
    if dl.add_cookie_to_db(number, host, cookie_str):
        return jsonify({'status': 'success', 'message': 'Cookie added/updated', 'id': number})
    return jsonify({'error': 'Failed to add cookie'}), 500

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000)