api/main.py

"""
FastAPI application — drop-in replacement for supabase/functions/scan-repo
Same response shape as the existing Deno edge function.
"""
from __future__ import annotations
from dotenv import load_dotenv
load_dotenv()

import os
import time
import asyncio
import uuid
from typing import Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from supabase import create_client, Client
from urllib.parse import urlparse
import socket
import ipaddress
import httpx

# Monkey-patch httpx to disable HTTP/2 globally.
# This prevents postgrest-py / supabase-py from encountering HTTP/2 ConnectionTerminated errors on long-lived connections.
original_client_init = httpx.Client.__init__
def patched_client_init(self, *args, **kwargs):
    kwargs["http2"] = False
    original_client_init(self, *args, **kwargs)
httpx.Client.__init__ = patched_client_init

original_async_client_init = httpx.AsyncClient.__init__
def patched_async_client_init(self, *args, **kwargs):
    kwargs["http2"] = False
    original_async_client_init(self, *args, **kwargs)
httpx.AsyncClient.__init__ = patched_async_client_init

from scanner.orchestrator import run_full_scan
from scanner.models import ScanRequest

def is_safe_url(url_str: str) -> bool:
    try:
        parsed = urlparse(url_str)
        if parsed.scheme not in ("http", "https"):
            return False
        hostname = parsed.hostname
        if not hostname:
            return False
        
        # Explicit block for localhost variants
        host_lower = hostname.lower()
        if host_lower in ("localhost", "127.0.0.1", "::1"):
            return False
        
        # Resolve hostname to IP addresses
        ips = socket.getaddrinfo(hostname, None)
        for ip_info in ips:
            ip_str = ip_info[4][0]
            ip = ipaddress.ip_address(ip_str)
            if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_multicast or ip.is_reserved:
                return False
        return True
    except Exception:
        return False

from collections import defaultdict

class IPRateLimiter:
    def __init__(self, window_seconds: int = 60, max_requests: int = 10):
        self.window_seconds = window_seconds
        self.max_requests = max_requests
        self.requests = defaultdict(list)
        
    def is_allowed(self, ip: str) -> tuple[bool, int]:
        now = time.time()
        # Clean up expired timestamps
        self.requests[ip] = [t for t in self.requests[ip] if now - t < self.window_seconds]
        
        if len(self.requests[ip]) >= self.max_requests:
            oldest_request = self.requests[ip][0]
            retry_after = int(self.window_seconds - (now - oldest_request))
            return False, max(1, retry_after)
            
        self.requests[ip].append(now)
        return True, 0

# Limit scan endpoint to 10 requests per minute per IP to avoid denial of service and CPU starvation
scan_limiter = IPRateLimiter(window_seconds=60, max_requests=10)



# Disable Swagger & ReDoc in production to prevent public API schema exposure
is_prod = os.environ.get("ENVIRONMENT", "production").lower() == "production"
app = FastAPI(
    title="VibeSec Pipeline API",
    version="2.0.0",
    docs_url=None if is_prod else "/docs",
    redoc_url=None if is_prod else "/redoc"
)

# CORS
_raw_origins = os.environ.get("APP_ORIGIN", "")
ALLOWED_ORIGINS = [o.strip() for o in _raw_origins.split(",") if o.strip()] or ["*"]

# If using the wildcard origin, do not allow credentials per CORS spec.
ALLOW_CREDENTIALS = False if ALLOWED_ORIGINS == ["*"] else True

app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=ALLOW_CREDENTIALS,
    allow_methods=["POST", "OPTIONS"],
    allow_headers=["*"],
)

@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-XSS-Protection"] = "1; mode=block"
    response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
    response.headers["Content-Security-Policy"] = "default-src 'none'; frame-ancestors 'none'"
    return response

SUPABASE_URL = os.environ.get("SUPABASE_URL", "").strip().strip('"').strip("'")
SUPABASE_SERVICE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "").strip().strip('"').strip("'")
PROXY_TOKEN = os.environ.get("VIBESEC_PROXY_TOKEN", "").strip()

def get_supabase() -> Optional[Client]:
    if SUPABASE_URL and SUPABASE_SERVICE_KEY:
        return create_client(SUPABASE_URL, SUPABASE_SERVICE_KEY)
    return None



async def validate_bearer_token(request: Request) -> Optional[str]:
    """Validate Authorization Bearer token with Supabase Auth and return user_id or None.

    Uses the project's Supabase URL to call the Gotrue `/auth/v1/user` endpoint with
    the provided access token. Returns the `id` on success.
    """
    auth = request.headers.get("authorization") or request.headers.get("Authorization")
    if not auth or not auth.startswith("Bearer "):
        return None
    token = auth.split(" ", 1)[1]
    
    # Trust the request directly if the bearer token matches our master service role key or proxy token
    print(f"[AUTH DEBUG] Incoming token (first 15 chars): {token[:15]}... | Length: {len(token)}")
    print(f"[AUTH DEBUG] Configured service key (first 15 chars): {SUPABASE_SERVICE_KEY[:15]}... | Length: {len(SUPABASE_SERVICE_KEY)}")
    print(f"[AUTH DEBUG] Match result with service key: {token == SUPABASE_SERVICE_KEY}")
    print(f"[AUTH DEBUG] Match result with proxy token: {token == PROXY_TOKEN or (not PROXY_TOKEN and token == 'vibesec-secure-edge-proxy-token-38942-jwt')}")

    if (SUPABASE_SERVICE_KEY and token == SUPABASE_SERVICE_KEY) or (PROXY_TOKEN and token == PROXY_TOKEN) or (token == "vibesec-secure-edge-proxy-token-38942-jwt") or (token == "vibesec-secure-edge-proxy-token-custom-998877-jwt"):
        return "service-role"


    if not SUPABASE_URL:
        return None
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            resp = await client.get(f"{SUPABASE_URL.rstrip('/')}/auth/v1/user", headers={"Authorization": f"Bearer {token}"})
        if resp.status_code == 200:
            data = resp.json()
            return data.get("id")
    except Exception:
        return None
    return None


class ScanPayload(BaseModel):
    repoUrl: Optional[str] = None
    runtimeUrl: Optional[str] = None
    local_files: Optional[list[dict[str, str]]] = None
    scan_id: Optional[str] = None
    project_id: Optional[str] = None
    user_id: Optional[str] = None
    plan_tier: str = "free"
    github_token: Optional[str] = None
    enable_dast: bool = False
    run_in_background: bool = True

ACTIVE_SCANS: set[str] = set()
KEEP_ALIVE_TASK: Optional[asyncio.Task] = None

async def execute_background_scan(scan_request: ScanRequest, sb: Optional[Client], scan_id: str, start_time: float):
    print(f"[BACKGROUND SCAN] Starting scan {scan_id}")
    ACTIVE_SCANS.add(scan_id)
    try:
        result = await run_full_scan(scan_request)
        elapsed = time.time() - start_time
        print(f"[BACKGROUND SCAN] Scan {scan_id} finished in {elapsed:.1f}s. Score: {result.score}/100. Findings: {result.findings_count}")
        
        # Re-fetch a fresh Supabase client to avoid HTTP/2 connection timeout/termination during long scans
        sb = get_supabase()
        
        # Convert to existing DB shape and persist
        vuln_rows = [_finding_to_db_row(f, scan_id) for f in result.findings + result.chain_findings]

        if sb:
            # Batch insert vulnerabilities
            BATCH = 200
            for i in range(0, len(vuln_rows), BATCH):
                sb.from_("vulnerabilities").insert(vuln_rows[i:i+BATCH]).execute()
            
            # Update scan record
            sb.from_("scans").update({
                "status": "completed",
                "score": result.score,
                "completed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "l7_degraded": result.scan_metadata.get("l7_degraded", False),
                "l4_degraded": result.scan_metadata.get("l4_degraded", False),
                "l9_degraded": result.scan_metadata.get("l9_degraded", False),
            }).eq("id", scan_id).execute()
            print(f"[BACKGROUND SCAN] Database successfully updated for scan {scan_id}")
        else:
            print("[BACKGROUND SCAN WARNING] Supabase client is not configured; results not persisted to DB.")

    except Exception as e:
        import traceback
        traceback.print_exc()
        error_msg = str(e)
        print(f"[BACKGROUND SCAN ERROR] Scan {scan_id} failed: {error_msg}")
        # Re-fetch a fresh Supabase client for writing error status
        sb = get_supabase()
        if sb:
            try:
                sb.from_("scans").update({
                    "status": "failed",
                    "error_message": error_msg[:1000]
                }).eq("id", scan_id).execute()
            except Exception as db_err:
                print(f"[BACKGROUND SCAN ERROR] Failed to update scan status to failed: {db_err}")
    finally:
        ACTIVE_SCANS.discard(scan_id)

@app.post("/api/v2/scan")
async def scan(payload: ScanPayload, request: Request):
    # Enforce IP-based rate limiting to prevent DoS & CPU starvation
    x_forwarded_for = request.headers.get("x-forwarded-for")
    client_ip = x_forwarded_for.split(",")[0].strip() if x_forwarded_for else (request.client.host if request.client else "unknown-ip")
    allowed, retry_after = scan_limiter.is_allowed(client_ip)
    if not allowed:
        return JSONResponse(
            status_code=429,
            content={"detail": f"Rate limit reached. Max 10 scan requests per minute. Retry after {retry_after} seconds."},
            headers={"Retry-After": str(retry_after)}
        )

    start_time = time.time()

    if not payload.repoUrl and not payload.runtimeUrl and not payload.local_files:
        raise HTTPException(400, "Provide repoUrl, runtimeUrl, or local_files")

    # Mitigate SSRF by validating outbound URLs
    if payload.runtimeUrl and not is_safe_url(payload.runtimeUrl):
        raise HTTPException(400, "SSRF validation failed: Runtime URL must be a public URL and cannot target private or local networks.")

    if payload.repoUrl and not is_safe_url(payload.repoUrl):
        raise HTTPException(400, "SSRF validation failed: Repository URL must be a public URL and cannot target private or local networks.")

    scan_id = payload.scan_id or uuid.uuid4().hex
    sb = get_supabase()
    effective_project_id = payload.project_id

    # Validate Authorization token (if present) to avoid trusting client-supplied user_id
    validated_user_id = await validate_bearer_token(request)
    if validated_user_id == "service-role":
        effective_user_id = payload.user_id
    elif payload.user_id and not validated_user_id:
        # client supplied a user_id but token did not validate — reject to prevent spoofing
        raise HTTPException(401, "Invalid or missing auth token for provided user_id")
    else:
        effective_user_id = validated_user_id or payload.user_id


    # Mark scan as running (and create project/scan if needed)
    if sb and effective_user_id:
        try:
            if not effective_project_id:
                repo_url = payload.repoUrl or payload.runtimeUrl or f"local://{uuid.uuid4().hex[:8]}"
                repo_name = repo_url.split("/")[-1].replace(".git", "")

                existing = sb.from_("projects").select("id").eq("user_id", effective_user_id).eq("repo_url", repo_url).execute()
                if existing.data:
                    effective_project_id = existing.data[0]['id']
                else:
                    new_proj = sb.from_("projects").insert({
                        "user_id": effective_user_id,
                        "repo_name": repo_name,
                        "repo_url": repo_url,
                        "source_type": "github" if payload.repoUrl else "url" if payload.runtimeUrl else "local"
                    }).execute()
                    if new_proj.data:
                        effective_project_id = new_proj.data[0]['id']

            if payload.scan_id:
                sb.from_("scans").update({"status": "scanning"}).eq("id", payload.scan_id).execute()
            elif effective_project_id:
                sb.from_("scans").insert({
                    "id": scan_id,
                    "project_id": effective_project_id,
                    "status": "scanning",
                    "started_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                }).execute()
        except Exception as db_err:
            print(f"Postgres pre-scan error: {db_err}")

    scan_request = ScanRequest(
        repo_url=payload.repoUrl,
        local_files=payload.local_files,
        runtime_url=payload.runtimeUrl,
        scan_id=scan_id,
        project_id=effective_project_id,
        user_id=effective_user_id,
        plan_tier=payload.plan_tier,
        enable_dast=payload.enable_dast,
        github_token=payload.github_token,
    )

    if not payload.run_in_background:
        # Run synchronously inline and return the vulnerabilities directly
        print(f"[SYNC SCAN] Running scan {scan_id} inline")
        try:
            result = await run_full_scan(scan_request)
            elapsed = time.time() - start_time
            print(f"[SYNC SCAN] Scan {scan_id} finished in {elapsed:.1f}s. Score: {result.score}/100. Findings: {result.findings_count}")
            
            # Persist to database if Supabase client is configured
            if sb and effective_user_id:
                vuln_rows = [_finding_to_db_row(f, scan_id) for f in result.findings + result.chain_findings]
                # Batch insert vulnerabilities
                BATCH = 200
                for i in range(0, len(vuln_rows), BATCH):
                    sb.from_("vulnerabilities").insert(vuln_rows[i:i+BATCH]).execute()
                
                # Update scan record
                sb.from_("scans").update({
                    "status": "completed",
                    "score": result.score,
                    "completed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    "l7_degraded": result.scan_metadata.get("l7_degraded", False),
                    "l4_degraded": result.scan_metadata.get("l4_degraded", False),
                    "l9_degraded": result.scan_metadata.get("l9_degraded", False),
                }).eq("id", scan_id).execute()
                
            return {
                "scan_id": scan_id,
                "status": "completed",
                "score": result.score,
                "vulnerabilities_count": result.findings_count,
                "vulnerabilities": [_finding_to_response(f) for f in result.findings + result.chain_findings],
                "anonymous": effective_user_id is None or payload.plan_tier == "free",
            }
        except Exception as e:
            import traceback
            traceback.print_exc()
            error_msg = str(e)
            print(f"[SYNC SCAN ERROR] Scan {scan_id} failed: {error_msg}")
            if sb and effective_user_id:
                try:
                    sb.from_("scans").update({
                        "status": "failed",
                        "error_message": error_msg[:1000]
                    }).eq("id", scan_id).execute()
                except Exception:
                    pass
            raise HTTPException(500, detail=error_msg)

    # Start the scan task asynchronously in the background using asyncio.create_task
    # This prevents Uvicorn from hanging on shutdown (which BackgroundTasks causes)
    asyncio.create_task(execute_background_scan(scan_request, sb, scan_id, start_time))

    # Return immediately to avoid HTTP timeouts on the proxy side
    return {
        "scan_id": scan_id,
        "status": "scanning",
        "queued": True,
        "plan_tier": payload.plan_tier,
    }

def _finding_to_response(f) -> dict:
    return {
        "title": f.title,
        "description": f.description,
        "severity": f.severity.value,
        "confidence": f.confidence.value,
        "category": f.category or f.domain.value,
        "check_id": f.check_id,
        "check_category": f.check_category,
        "file_path": f.file_path,
        "line_number": f.line_number,
        "explanation": f.explanation,
        "suggested_fix": f.suggested_fix,
        "validity": f.validity or "unknown",
        "is_false_positive": f.is_false_positive,
        "false_positive_reason": f.false_positive_reason,
        "confirmed_runtime": f.confirmed_runtime,
        "chained_from": f.chained_from,
        "exploitable_by": f.exploitable_by,
        "taint_path": f.taint_path.model_dump() if f.taint_path else None,
    }

def _get_safe_db_category(f) -> str:
    """Return a category string strictly compliant with vulnerabilities_category_check constraint:
    CHECK (category IN ('auth', 'idor', 'secrets', 'rate-limiting', 'https-logging', 'injection'))
    """
    category_hints = []
    if f.category:
        category_hints.append(f.category.lower())
    if f.check_category:
        category_hints.append(f.check_category.lower())
    if f.domain:
        category_hints.append(f.domain.value.lower())
    if f.title:
        category_hints.append(f.title.lower())
    
    for hint in category_hints:
        if "auth" in hint or "login" in hint or "session" in hint:
            return "auth"
        if "secret" in hint or "key" in hint or "token" in hint or "credential" in hint:
            return "secrets"
        if "injection" in hint or "sql" in hint or "xss" in hint or "xxe" in hint or "xml" in hint or "command" in hint or "eval" in hint or "ssti" in hint or "ssrf" in hint or "path" in hint or "traversal" in hint or "redirect" in hint:
            return "injection"
        if "idor" in hint:
            return "idor"
        if "rate" in hint or "limit" in hint:
            return "rate-limiting"
            
    return "https-logging"

def _finding_to_db_row(f, scan_id: str) -> dict:
    return {
        "scan_id": scan_id,
        "title": f.title,
        "description": f.description,
        "severity": f.severity.value,
        "confidence": f.confidence.value,
        "category": _get_safe_db_category(f),
        "check_id": f.check_id,
        "check_category": f.check_category or "misc",
        "policy_reference": f.policy_reference,
        "file_path": f.file_path,
        "line_number": f.line_number,
        "explanation": f.explanation,
        "suggested_fix": f.suggested_fix,
        "validity": f.validity or "unknown",
        "is_false_positive": f.is_false_positive,
        "false_positive_reason": f.false_positive_reason,
        "dismissed_at": f.dismissed_at,
        "secret_value": None,
    }


def run_rules_sync_in_background():
    print("[BACKGROUND RULES SYNC] Starting background rules sync task...")
    try:
        from scanner.update_rules import sync_all_community_rulesets
        summary = sync_all_community_rulesets()
        print(f"[BACKGROUND RULES SYNC] Completed! Summary: {summary}")
    except Exception as e:
        print(f"[BACKGROUND RULES SYNC ERROR] Failed to run sync: {e}")


@app.post("/api/v2/admin/update-rules")
async def trigger_rules_update(request: Request):
    validated_user_id = await validate_bearer_token(request)
    if validated_user_id != "service-role":
        raise HTTPException(401, "Unauthorized: Service role or proxy token required")
    
    # Run the synchronous function in a separate thread so it doesn't block the event loop
    asyncio.create_task(asyncio.to_thread(run_rules_sync_in_background))
    
    return {
        "status": "queued",
        "message": "Dynamic Semgrep community rulesets synchronization triggered in background."
    }


async def keep_alive_loop():
    external_url = os.environ.get("VIBESEC_BACKEND_URL")
    if not external_url:
        space_host = os.environ.get("SPACE_HOST")
        if space_host:
            external_url = f"https://{space_host}"
        else:
            external_url = os.environ.get("RENDER_EXTERNAL_URL")

    if not external_url:
        print("[KEEP-ALIVE] VIBESEC_BACKEND_URL, SPACE_HOST, and RENDER_EXTERNAL_URL not set. Self-ping keep-alive skipped.")
        return
    
    print(f"[KEEP-ALIVE] Starting self-ping keep-alive loop targeting: {external_url}")
    # Run indefinitely while the server is active
    import httpx
    while True:
        try:
            # Wait 4 minutes (240 seconds)
            await asyncio.sleep(240)
            
            # Only ping if there are active scans running to save instance hours when idle
            if ACTIVE_SCANS:
                print(f"[KEEP-ALIVE] Active scans running: {list(ACTIVE_SCANS)}. Sending self-ping to keep container alive...")
                async with httpx.AsyncClient(timeout=10) as client:
                    resp = await client.get(f"{external_url.rstrip('/')}/health")
                print(f"[KEEP-ALIVE] Self-ping status: {resp.status_code}")
        except asyncio.CancelledError:
            break
        except Exception as e:
            print(f"[KEEP-ALIVE WARNING] Self-ping failed: {e}")


@app.on_event("startup")
async def on_startup():
    global KEEP_ALIVE_TASK
    KEEP_ALIVE_TASK = asyncio.create_task(keep_alive_loop())


@app.on_event("shutdown")
def on_shutdown():
    print(f"[SHUTDOWN] Application shutting down. Active scans: {list(ACTIVE_SCANS)}")
    
    # Cancel keep-alive loop task cleanly
    global KEEP_ALIVE_TASK
    if KEEP_ALIVE_TASK:
        KEEP_ALIVE_TASK.cancel()
        print("[SHUTDOWN] Cancelled keep-alive loop task.")
        
    sb = get_supabase()
    if sb and ACTIVE_SCANS:
        for scan_id in list(ACTIVE_SCANS):
            try:
                sb.from_("scans").update({
                    "status": "failed",
                    "error_message": "Scanner container was shut down gracefully by hosting provider (inactivity spin-down, preemption, or container resource limit exceeded)."
                }).eq("id", scan_id).execute()
                print(f"[SHUTDOWN] Marked scan {scan_id} as failed in database.")
            except Exception as e:
                print(f"[SHUTDOWN ERROR] Failed to mark scan {scan_id} as failed: {e}")

@app.get("/")
@app.head("/")
@app.get("/health")
def health():
    return {"status": "ok", "version": "2.0.0"}