statdx/scrapers/cache_server.py

#!/usr/bin/env python3
import os
import json
import base64
import hashlib
from pathlib import Path
from typing import Dict, Optional
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# Import exact helpers from capture_passive_playwright_async
from scrapers.capture_passive_playwright_async import (
    sanitize,
    guess_ext,
    now_ts,
    try_pretty_json,
    save_binary,
    save_text
)

app = FastAPI(title="STATdx Central Cache Server")

# Allow CORS for options requests and standard requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

OUTPUT_DIR = Path("xhr_captured_async")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
INDEX_PATH = OUTPUT_DIR / "capture_index.jsonl"

# Persistent in-memory index
in_memory_index: Dict[str, list] = {}

def load_existing_index():
    if INDEX_PATH.exists():
        print(f"Loading existing index from {INDEX_PATH}...")
        try:
            with open(INDEX_PATH, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        entry = json.loads(line)
                        ch = entry.get("content_hash")
                        if ch:
                            in_memory_index.setdefault(ch, []).append(entry)
                    except Exception:
                        pass
            print(f"Loaded {len(in_memory_index)} unique hashes into memory.")
        except Exception as e:
            print(f"Error loading index: {e}")

# Load index at import/startup
load_existing_index()

class HashCheckRequest(BaseModel):
    content_hash: str

class CapturePayload(BaseModel):
    url: str
    resource_type: str
    status: int
    timestamp: Optional[str] = None
    headers: Dict[str, str]
    body_base64: str
    content_type: str

@app.post("/api/check-hash")
async def check_hash(payload: HashCheckRequest):
    exists = payload.content_hash in in_memory_index
    return {"exists": exists}

@app.post("/api/capture")
async def capture(payload: CapturePayload):
    try:
        # Decode body
        body_bytes = base64.b64decode(payload.body_base64)
        content_hash = hashlib.sha256(body_bytes).hexdigest()
        
        # Check deduplication
        if content_hash in in_memory_index:
            # We already have this exact content saved
            # Just return and do not write duplicate files
            return {"status": "skipped", "reason": "content_hash already cached"}

        url = payload.url
        ts = payload.timestamp or now_ts()
        safe = sanitize(url.replace('https://', '').replace('http://', ''))
        h = hashlib.sha1((url + ts).encode('utf-8')).hexdigest()[:8]

        body_path = None
        excerpt = None
        ctype = payload.content_type.lower()
        rtype = payload.resource_type.lower()

        # Handle text content
        if 'application/json' in ctype or ctype.startswith('text/') or 'javascript' in ctype:
            try:
                txt = body_bytes.decode('utf-8')
            except Exception:
                txt = body_bytes.decode('latin-1', errors='replace')

            pretty = try_pretty_json(txt)
            if pretty is not None:
                use_text = pretty
                excerpt = pretty[:800]
                ext = 'json'
            else:
                use_text = txt
                excerpt = txt[:800]
                ext = 'txt'

            content_sha1 = hashlib.sha1(use_text.encode('utf-8')).hexdigest()
            fname = f"{safe}_{content_sha1}.{ext}"
            body_path = str(OUTPUT_DIR / fname)

            if not os.path.exists(body_path):
                save_text(OUTPUT_DIR / fname, use_text)
        else:
            # Handle binary content
            ext = guess_ext(ctype)
            content_sha1 = hashlib.sha1(body_bytes).hexdigest()
            fname = f"{safe}_{content_sha1}.{ext}"

            if ctype.startswith('image/'):
                subdir = OUTPUT_DIR / 'images'
                subdir.mkdir(parents=True, exist_ok=True)
                body_path = str(subdir / fname)
                if not os.path.exists(body_path):
                    save_binary(subdir / fname, body_bytes)
            else:
                subdir = OUTPUT_DIR / 'assets'
                subdir.mkdir(parents=True, exist_ok=True)
                body_path = str(subdir / fname)
                if not os.path.exists(body_path):
                    save_binary(subdir / fname, body_bytes)

        # Meta file creation
        meta_base = content_hash or h
        new_meta_name = f"{safe}_{meta_base}.meta.json"
        new_meta_path = OUTPUT_DIR / new_meta_name

        meta = {
            'url': url,
            'resource_type': rtype,
            'status': payload.status,
            'timestamp': ts,
            'response_headers': payload.headers,
            'response_body_file': body_path,
            'response_excerpt': excerpt,
            'content_hash': content_hash,
        }

        save_text(new_meta_path, json.dumps(meta, ensure_ascii=False, indent=2))

        # Write to log index
        entry = {
            'url': url,
            'resource_type': rtype,
            'timestamp': ts,
            'body_file': body_path,
            'meta_file': str(new_meta_path),
            'excerpt': excerpt,
            'content_hash': content_hash
        }

        with open(INDEX_PATH, 'a', encoding='utf-8') as idx:
            idx.write(json.dumps(entry, ensure_ascii=False) + '\n')

        # Update in-memory index
        in_memory_index.setdefault(content_hash, []).append(entry)

        print(f"SAVED: {url} -> {body_path}")
        return {"status": "saved", "path": body_path}

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("scrapers.cache_server:app", host="0.0.0.0", port=8000, reload=True)