speedster/__init__.py

import os
import sys
import asyncio
import aiohttp
import urllib.parse
import argparse
import shutil
import random
import hashlib
from tqdm.asyncio import tqdm
import nest_asyncio

# Apply nest_asyncio globally
nest_asyncio.apply()

# --- STEALTH AGENTS ---
USER_AGENTS = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.2 Safari/605.1.15",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:123.0) Gecko/20100101 Firefox/123.0",
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
]

class Speedster:
    def __init__(self, num_threads=16, chunk_size_mb=2, max_retries=5):
        self.num_threads = num_threads
        self.chunk_size = chunk_size_mb * 1024 * 1024 
        self.max_retries = max_retries
        self.is_windows = os.name == 'nt'

    def _get_headers(self, token=None):
        # Stealth Mode: Pick random agent
        headers = {
            "User-Agent": random.choice(USER_AGENTS),
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
            "Accept-Language": "en-US,en;q=0.5",
        }
        # Priority: Function Arg > Env Var
        final_token = token or os.environ.get("CIVITAI_TOKEN") or os.environ.get("HF_TOKEN")
        if final_token:
            headers["Authorization"] = f"Bearer {final_token}"
        return headers

    def _write_data(self, fd, data, offset):
        if self.is_windows:
            os.lseek(fd, offset, 0)
            os.write(fd, data)
        else:
            os.pwrite(fd, data, offset)

    def _check_disk_space(self, path, required_bytes):
        total, used, free = shutil.disk_usage(os.path.dirname(os.path.abspath(path)) or '.')
        if free < required_bytes:
            raise OSError(f"❌ Disk Full! Required: {required_bytes/1024**3:.2f}GB, Free: {free/1024**3:.2f}GB")

    async def _fetch_chunk(self, session, url, start, end, fd, headers, pbar):
        chunk_headers = headers.copy()
        chunk_headers['Range'] = f'bytes={start}-{end}'
        
        for attempt in range(self.max_retries):
            try:
                async with session.get(url, headers=chunk_headers, timeout=60) as response:
                    response.raise_for_status()
                    offset = start
                    async for chunk in response.content.iter_chunked(self.chunk_size):
                        if chunk:
                            self._write_data(fd, chunk, offset)
                            offset += len(chunk)
                            pbar.update(len(chunk))
                return
            except Exception as e:
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(1 * (attempt + 1))
                else:
                    # tqdm.write(f"❌ Chunk failed: {e}") # Silent fail to reduce spam
                    raise e

    async def _download_logic(self, url, dest_path, token=None, extract=False, verify=False):
        headers = self._get_headers(token)
        
        async with aiohttp.ClientSession() as session:
            # 1. Metadata & Smart Filename
            try:
                async with session.get(url, headers=headers, allow_redirects=True) as response:
                    final_url = response.url
                    total_size = int(response.headers.get('Content-Length', 0))
                    
                    filename = "downloaded_file.bin"
                    if os.path.isdir(dest_path):
                        if 'Content-Disposition' in response.headers:
                            try:
                                cd = response.headers['Content-Disposition']
                                filename = cd.split('filename=')[1].strip('"')
                            except:
                                filename = os.path.basename(urllib.parse.urlparse(str(final_url)).path)
                        else:
                            path_filename = os.path.basename(urllib.parse.urlparse(str(final_url)).path)
                            if path_filename: filename = path_filename
                        
                        filename = filename.split('?')[0]
                        final_path = os.path.join(dest_path, filename)
                    else:
                        final_path = dest_path

                # 2. Smart Skip (File Exists & Size Match)
                if os.path.exists(final_path):
                    local_size = os.path.getsize(final_path)
                    if local_size == total_size and total_size > 0:
                        print(f"⏩ [Skip] {os.path.basename(final_path)} exists & valid.")
                        return final_path

                # 3. Disk Check
                self._check_disk_space(dest_path, total_size)

                print(f"🚀 [Speedster] {os.path.basename(final_path)}")
                print(f"📦 Size: {total_size / (1024**3):.2f} GB | Threads: {self.num_threads}")

                # 4. Atomic Download (.part file)
                part_path = final_path + ".part"
                fd = os.open(part_path, os.O_RDWR | os.O_CREAT | os.O_BINARY if self.is_windows else os.O_RDWR | os.O_CREAT)
                os.ftruncate(fd, total_size)

                chunk_size = total_size // self.num_threads
                chunks = []
                for i in range(self.num_threads):
                    start = i * chunk_size
                    end = total_size - 1 if i == self.num_threads - 1 else (start + chunk_size - 1)
                    chunks.append((start, end))

                pbar = tqdm(total=total_size, unit='iB', unit_scale=True, unit_divisor=1024, desc="⚡ DL")
                tasks = [asyncio.create_task(self._fetch_chunk(session, final_url, s, e, fd, headers, pbar)) for s, e in chunks]
                
                await asyncio.gather(*tasks)
                
                pbar.close()
                os.close(fd)
                
                # 5. Finalize (Rename .part -> .safetensors)
                if os.path.exists(final_path):
                    os.remove(final_path)
                os.rename(part_path, final_path)
                print("✅ Download Complete!")

                # 6. Auto-Extract
                if extract:
                    print(f"📂 Extracting {filename}...")
                    try:
                        shutil.unpack_archive(final_path, dest_path)
                        print("✅ Extracted successfully.")
                    except Exception as e:
                        print(f"⚠️ Extraction failed: {e}")

                return final_path

            except Exception as e:
                print(f"❌ Error: {e}")
                if 'part_path' in locals() and os.path.exists(part_path):
                    os.remove(part_path) # Cleanup corrupt partial
                return None

    def download(self, url, dest_path=".", token=None, extract=False, verify=False):
        loop = asyncio.get_event_loop()
        if loop.is_running():
            nest_asyncio.apply()
            return loop.run_until_complete(self._download_logic(url, dest_path, token, extract, verify))
        else:
            return asyncio.run(self._download_logic(url, dest_path, token, extract, verify))

# --- CLI HANDLER ---
def cli_main():
    parser = argparse.ArgumentParser(description="Speedster v16.0.3 - Juggernaut Update")
    parser.add_argument("input", help="URL file OR path to .txt file for batch download")
    parser.add_argument("--out", "-o", default=".", help="Output folder")
    parser.add_argument("--token", "-t", help="Auth Token (Optional if env var set)")
    parser.add_argument("--threads", type=int, default=16, help="Thread count")
    parser.add_argument("--extract", "-x", action="store_true", help="Auto extract archives")
    parser.add_argument("--batch", "-b", action="store_true", help="Force treat input as batch file")
    
    args = parser.parse_args()
    engine = Speedster(num_threads=args.threads)
    
    # BATCH MODE LOGIC
    urls = []
    if args.batch or (os.path.isfile(args.input) and args.input.endswith(".txt")):
        print(f"📜 Batch Mode: Reading {args.input}")
        with open(args.input, 'r') as f:
            urls = [line.strip() for line in f if line.strip() and not line.startswith("#")]
    else:
        urls = [args.input]
        
    for i, url in enumerate(urls):
        if len(urls) > 1: print(f"\n--- Processing {i+1}/{len(urls)} ---")
        engine.download(url, args.out, args.token, extract=args.extract)

if __name__ == "__main__":
    cli_main()