Upload app.py with huggingface_hub

app.py

@@ -1,101 +1,253 @@
-"""Cloudflare WARP SOCKS5 Proxy Server with HTTP CONNECT forwarding
+"""
+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.
 
-Exposes WARP as an HTTP proxy that other HF Spaces can use.
-Routes: /proxy?url=<url> - fetches URL through WARP VPN
-        /connect - HTTP CONNECT tunnel (for YouTube API)
+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 asyncio
+import time
+from pathlib import Path
+
 import httpx
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Query, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
 app = FastAPI(title="WARP SOCKS5 Proxy", version="2.0")
 
-WARP_CONNECTED = False
+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():
-    global WARP_CONNECTED
-    print("[WARP-PROXY] Starting WARP daemon...")
-    subprocess.Popen(["warp-svc"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-    asyncio.create_task(_connect_warp())
-
-async def _connect_warp():
-    global WARP_CONNECTED
-    await asyncio.sleep(3)
-    for attempt in range(5):
-        try:
-            subprocess.run(["warp-cli", "--accept-tos", "registration", "new"], capture_output=True, timeout=10)
-            subprocess.run(["warp-cli", "--accept-tos", "mode", "proxy"], capture_output=True, timeout=10)
-            subprocess.run(["warp-cli", "--accept-tos", "connect"], capture_output=True, timeout=10)
-            for i in range(30):
-                result = subprocess.run(["warp-cli", "--accept-tos", "status"], capture_output=True, text=True, timeout=5)
-                if "connected" in result.stdout.lower():
-                    print(f"[WARP-PROXY] CONNECTED! VPN active.")
-                    WARP_CONNECTED = True
-                    return
-                await asyncio.sleep(2)
-        except Exception as e:
-            print(f"[WARP-PROXY] Attempt {attempt+1} error: {e}")
-            await asyncio.sleep(3)
-    print("[WARP-PROXY] Failed - running without VPN")
+    """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():
-    return {"status": "ok" if WARP_CONNECTED else "connecting", "warp_connected": WARP_CONNECTED}
+    """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("/health")
-async def health():
-    return {"warp_connected": WARP_CONNECTED}
 
 @app.get("/proxy")
-async def proxy_get(url: str = ""):
-    """Fetch a URL through the WARP VPN SOCKS5 proxy."""
-    if not url:
-        return JSONResponse({"error": "url parameter required"}, 400)
+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="socks5://127.0.0.1:40000" if WARP_CONNECTED else None,
-            timeout=30,
-            follow_redirects=True
+            proxy=WARP_SOCKS5,
+            follow_redirects=follow_redirects,
+            timeout=timeout,
         ) as client:
-            resp = await client.get(url)
+            # 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)
+                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": str(e)}, 500)
+        return JSONResponse(
+            {"error": f"Proxy error: {e}"},
+            status_code=500
+        )
 
-@app.post("/proxy")
-async def proxy_post(request: Request):
-    """POST a request through the WARP VPN SOCKS5 proxy."""
+
+@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:
-        body = await request.body()
-        target_url = request.headers.get("X-Target-URL", "")
-        if not target_url:
-            return JSONResponse({"error": "X-Target-URL header required"}, 400)
-        
-        headers = {}
-        for key, value in request.headers.items():
-            if key.lower() not in ('host', 'x-target-url', 'content-length'):
-                headers[key] = value
-        
         async with httpx.AsyncClient(
-            proxy="socks5://127.0.0.1:40000" if WARP_CONNECTED else None,
-            timeout=60,
+            proxy=WARP_SOCKS5,
+            follow_redirects=True,
+            timeout=120.0,
         ) as client:
-            resp = await client.post(target_url, content=body, headers=headers)
-            return Response(
-                content=resp.content,
-                status_code=resp.status_code,
-                headers=dict(resp.headers)
-            )
+            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": str(e)}, 500)
+        return JSONResponse(
+            {"error": f"Download failed: {e}"},
+            status_code=500
+        )
+
 
 if __name__ == "__main__":
     import uvicorn