app.py

"""
WARP SOCKS5 Proxy Space - Cloudflare WARP VPN in Docker

This Space runs Cloudflare WARP as a SOCKS5 proxy on port 40000.
It also provides an HTTP forwarding endpoint so other HF Spaces can
route requests through WARP without needing SOCKS5 support.

Endpoints:
  GET  /          - Health check + WARP status
  GET  /ping      - Simple health check
  GET  /status    - WARP connection status
  GET  /proxy     - HTTP proxy: ?url=<target_url>  (forwards request through WARP)
  POST /proxy     - HTTP proxy: body={"url": "...", "method": "GET", "headers": {}}
  GET  /download  - Download a file through WARP: ?url=<file_url>&output=json
"""

import asyncio
import os
import subprocess
import time
from pathlib import Path

import httpx
from fastapi import FastAPI, Query, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse

app = FastAPI(title="WARP SOCKS5 Proxy", version="2.0")

WARP_SOCKS5 = "socks5://127.0.0.1:40000"
_warp_ready = False
_warp_start_time = 0


def is_warp_connected() -> bool:
    """Check if WARP is connected."""
    try:
        result = subprocess.run(
            ["warp-cli", "--accept-tos", "status"],
            capture_output=True, text=True, timeout=5
        )
        return "Connected" in result.stdout
    except Exception:
        return False


async def wait_for_warp(max_seconds: int = 30):
    """Wait for WARP to connect."""
    global _warp_ready
    for i in range(max_seconds):
        if is_warp_connected():
            _warp_ready = True
            print(f"[WARP] Connected after {i+1}s")
            return True
        await asyncio.sleep(1)
    print(f"[WARP] Failed to connect after {max_seconds}s")
    return False


@app.on_event("startup")
async def startup():
    """Start WARP daemon and connect."""
    global _warp_start_time
    _warp_start_time = time.time()
    
    print("[STARTUP] WARP Proxy Space starting...")
    
    # warp-svc should already be running from entrypoint.sh
    # Just need to register and connect
    print("[STARTUP] Registering WARP...")
    subprocess.run(["warp-cli", "--accept-tos", "registration", "new"],
                   capture_output=True, timeout=10)
    
    print("[STARTUP] Setting proxy mode...")
    subprocess.run(["warp-cli", "--accept-tos", "mode", "proxy"],
                   capture_output=True, timeout=10)
    
    print("[STARTUP] Connecting WARP...")
    subprocess.run(["warp-cli", "--accept-tos", "connect"],
                   capture_output=True, timeout=10)
    
    # Wait for connection in background
    asyncio.create_task(wait_for_warp(60))
    
    print("[STARTUP] WARP Proxy Space ready!")


@app.get("/")
async def root():
    """Health check with WARP status."""
    connected = is_warp_connected()
    uptime = int(time.time() - _warp_start_time) if _warp_start_time else 0
    return {
        "status": "ok" if connected else "starting",
        "warp_connected": connected,
        "socks5_proxy": WARP_SOCKS5 if connected else None,
        "uptime_seconds": uptime,
        "usage": {
            "http_proxy": "/proxy?url=<target_url>",
            "download": "/download?url=<file_url>",
            "socks5": WARP_SOCKS5,
        }
    }


@app.get("/ping")
async def ping():
    """Simple health check."""
    connected = is_warp_connected()
    return {"status": "alive", "warp_connected": connected}


@app.get("/status")
async def status():
    """Get WARP connection status."""
    try:
        result = subprocess.run(
            ["warp-cli", "--accept-tos", "status"],
            capture_output=True, text=True, timeout=5
        )
        return {"raw": result.stdout.strip(), "connected": "Connected" in result.stdout}
    except Exception as e:
        return {"raw": str(e), "connected": False}


@app.get("/proxy")
async def proxy_get(url: str = Query(..., description="Target URL to proxy")):
    """Forward a GET request through WARP SOCKS5 proxy."""
    return await _forward_request(url, "GET")


@app.post("/proxy")
async def proxy_post(request: Request):
    """Forward a POST request through WARP SOCKS5 proxy.
    
    Body: {"url": "...", "method": "GET", "headers": {}, "body": "..."}
    """
    try:
        body = await request.json()
        target_url = body.get("url", "")
        method = body.get("method", "GET").upper()
        headers = body.get("headers", {})
        req_body = body.get("body")
        return await _forward_request(target_url, method, headers, req_body)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=400)


async def _forward_request(url: str, method: str = "GET",
                           headers: dict = None, body=None,
                           follow_redirects: bool = True,
                           timeout: float = 60.0):
    """Forward a request through the WARP SOCKS5 proxy."""
    if not _warp_ready and not is_warp_connected():
        # Try to wait a bit
        connected = await wait_for_warp(10)
        if not connected:
            return JSONResponse(
                {"error": "WARP not connected", "status": "starting"},
                status_code=503
            )
    
    try:
        async with httpx.AsyncClient(
            proxy=WARP_SOCKS5,
            follow_redirects=follow_redirects,
            timeout=timeout,
        ) as client:
            # Filter out hop-by-hop headers
            safe_headers = {}
            if headers:
                skip = {"host", "connection", "keep-alive", "transfer-encoding",
                        "te", "trailer", "upgrade", "proxy-authorization"}
                for k, v in headers.items():
                    if k.lower() not in skip:
                        safe_headers[k] = v
            
            resp = await client.request(
                method=method,
                url=url,
                headers=safe_headers,
                content=body if isinstance(body, bytes) else (body.encode() if isinstance(body, str) else None),
            )
            
            # Return the response
            return Response(
                content=resp.content,
                status_code=resp.status_code,
                headers=dict(resp.headers),
            )
    except httpx.ConnectError as e:
        return JSONResponse(
            {"error": f"WARP proxy connection failed: {e}"},
            status_code=502
        )
    except httpx.TimeoutException:
        return JSONResponse(
            {"error": "Request timed out through WARP proxy"},
            status_code=504
        )
    except Exception as e:
        return JSONResponse(
            {"error": f"Proxy error: {e}"},
            status_code=500
        )


@app.get("/download")
async def download_file(url: str = Query(..., description="File URL to download through WARP")):
    """Download a file through WARP and stream it back.
    
    Query params:
      url: The file URL to download
      output: "stream" (default) or "json" (returns base64)
    """
    if not _warp_ready and not is_warp_connected():
        connected = await wait_for_warp(10)
        if not connected:
            return JSONResponse(
                {"error": "WARP not connected"},
                status_code=503
            )
    
    try:
        async with httpx.AsyncClient(
            proxy=WARP_SOCKS5,
            follow_redirects=True,
            timeout=120.0,
        ) as client:
            async with client.stream("GET", url) as resp:
                if resp.status_code != 200:
                    return JSONResponse(
                        {"error": f"Upstream returned {resp.status_code}"},
                        status_code=resp.status_code
                    )
                
                content_type = resp.headers.get("content-type", "application/octet-stream")
                
                return StreamingResponse(
                    resp.aiter_bytes(chunk_size=65536),
                    media_type=content_type,
                    headers={
                        "Content-Disposition": f'attachment; filename="{url.split("/")[-1].split("?")[0]}"',
                    }
                )
    except Exception as e:
        return JSONResponse(
            {"error": f"Download failed: {e}"},
            status_code=500
        )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)