# -*- coding: utf-8 -*-
"""
OK ITHAKI Private Relay Node V249436
====================================
Application-specific fixed-destination HTTPS reverse relay.

This node is NOT an open proxy and does NOT claim full Tor equivalence.
It forwards only to one configured upstream (another ITHAKI relay node or
https://ok-ithaki.gr). For meaningful separation, deploy relay nodes on
independent hosts/providers.

Environment:
  ITHAKI_RELAY_UPSTREAM=https://ok-ithaki.gr
  ITHAKI_RELAY_ACCESS_KEY=long-random-private-code
  ITHAKI_RELAY_PUBLIC_BASE=https://relay1.example.gr
  ITHAKI_RELAY_HOP_NAME=relay-1
  ITHAKI_RELAY_MAX_BODY_MB=300

Run:
  uvicorn ithaki_private_relay_v249436:app --host 127.0.0.1 --port 8021
"""
from __future__ import annotations

import hashlib
import hmac
import os
import time
from typing import Iterable
from urllib.parse import quote, unquote

import httpx
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import Response, StreamingResponse, JSONResponse

VERSION = "V249436-ithaki-private-relay-node"
UPSTREAM = os.getenv("ITHAKI_RELAY_UPSTREAM", "https://ok-ithaki.gr").rstrip("/")
ACCESS_KEY = os.getenv("ITHAKI_RELAY_ACCESS_KEY", "").strip()
PUBLIC_BASE = os.getenv("ITHAKI_RELAY_PUBLIC_BASE", "").rstrip("/")
HOP_NAME = os.getenv("ITHAKI_RELAY_HOP_NAME", "ithaki-relay").strip() or "ithaki-relay"
MAX_BODY = max(1, min(int(os.getenv("ITHAKI_RELAY_MAX_BODY_MB", "300")), 1024)) * 1024 * 1024

app = FastAPI(title="OK ITHAKI Private Relay", version=VERSION)
client = httpx.AsyncClient(follow_redirects=False, timeout=httpx.Timeout(90.0, connect=15.0))

HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade", "host", "content-length",
}
PRIVATE_HEADERS = {
    "x-forwarded-for", "x-real-ip", "forwarded", "cf-connecting-ip", "true-client-ip",
    "x-client-ip", "x-cluster-client-ip", "fastly-client-ip",
}


def _secure_equal(a: str, b: str) -> bool:
    return hmac.compare_digest((a or "").encode("utf-8"), (b or "").encode("utf-8"))


def _access_value(request: Request) -> str:
    raw = request.cookies.get("ithaki_relay_access", "")
    if raw:
        try:
            return unquote(raw)
        except Exception:
            return raw
    return request.headers.get("x-ithaki-relay-access", "")


def _authorized(request: Request) -> bool:
    return not ACCESS_KEY or _secure_equal(_access_value(request), ACCESS_KEY)


def _target(path: str, query: str = "") -> str:
    path = "/" + path.lstrip("/")
    return UPSTREAM + path + (("?" + query) if query else "")


def _request_headers(request: Request) -> dict[str, str]:
    out: dict[str, str] = {}
    for k, v in request.headers.items():
        lk = k.lower()
        if lk in HOP_BY_HOP or lk in PRIVATE_HEADERS or lk == "cookie":
            continue
        out[k] = v
    # Preserve application cookies, but never forward the relay access credential.
    cookies = [f"{k}={v}" for k, v in request.cookies.items() if k != "ithaki_relay_access"]
    if cookies:
        out["cookie"] = "; ".join(cookies)
    previous = request.headers.get("x-ithaki-relay-chain", "").strip()
    out["x-ithaki-relay-chain"] = (previous + "," if previous else "") + HOP_NAME
    out["x-ithaki-relay-version"] = VERSION
    return out


def _response_headers(headers: Iterable[tuple[str, str]]) -> dict[str, str]:
    out: dict[str, str] = {}
    for k, v in headers:
        lk = k.lower()
        if lk in HOP_BY_HOP or lk in {"content-encoding", "content-length"}:
            continue
        if lk in {"server", "via", "x-powered-by"}:
            continue
        if lk == "location" and PUBLIC_BASE and v.startswith(UPSTREAM):
            v = PUBLIC_BASE + v[len(UPSTREAM):]
        out[k] = v
    out["x-ithaki-private-relay"] = HOP_NAME
    out["x-ithaki-private-relay-truth"] = "app-specific-not-full-tor"
    return out


@app.get("/relay/v1/health")
async def relay_health(request: Request):
    if not _authorized(request):
        raise HTTPException(status_code=401, detail="relay_access_denied")
    upstream_hash = hashlib.sha256(UPSTREAM.encode("utf-8")).hexdigest()[:12]
    return {
        "ok": True,
        "version": VERSION,
        "hop": HOP_NAME,
        "upstream_hash": upstream_hash,
        "fixed_destination_only": True,
        "open_proxy": False,
        "full_tor_equivalent": False,
        "independent_nodes_required": True,
    }


@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def relay_http(path: str, request: Request):
    if path.startswith("relay/v1/"):
        raise HTTPException(status_code=404, detail="relay_control_not_found")
    if not _authorized(request):
        return JSONResponse({"ok": False, "reason": "relay_access_denied"}, status_code=401)
    body = await request.body()
    if len(body) > MAX_BODY:
        return JSONResponse({"ok": False, "reason": "body_too_large"}, status_code=413)
    target = _target(path, request.url.query)
    try:
        upstream = await client.request(
            request.method,
            target,
            headers=_request_headers(request),
            content=body,
        )
    except httpx.RequestError:
        return JSONResponse({"ok": False, "reason": "relay_upstream_unreachable"}, status_code=502)
    content = upstream.content
    ctype = upstream.headers.get("content-type", "")
    # Rewrite hardcoded origin references in text assets so subrequests remain on the relay origin.
    if PUBLIC_BASE and any(x in ctype for x in ("text/html", "javascript", "text/css", "application/json")):
        try:
            content = content.decode(upstream.encoding or "utf-8").replace(UPSTREAM, PUBLIC_BASE).encode("utf-8")
        except Exception:
            pass
    return Response(content=content, status_code=upstream.status_code, headers=_response_headers(upstream.headers.multi_items()), media_type=None)


@app.websocket("/{path:path}")
async def relay_websocket(websocket: WebSocket, path: str):
    code = websocket.cookies.get("ithaki_relay_access", "") or websocket.headers.get("x-ithaki-relay-access", "")
    try:
        code = unquote(code)
    except Exception:
        pass
    if ACCESS_KEY and not _secure_equal(code, ACCESS_KEY):
        await websocket.close(code=4401)
        return
    try:
        import websockets
    except Exception:
        await websocket.close(code=1011, reason="websockets_dependency_missing")
        return
    qs = websocket.url.query
    upstream_http = _target(path, qs)
    upstream_ws = upstream_http.replace("https://", "wss://", 1).replace("http://", "ws://", 1)
    await websocket.accept()
    try:
        async with websockets.connect(upstream_ws, max_size=64 * 1024 * 1024, open_timeout=15) as remote:
            async def client_to_remote():
                while True:
                    msg = await websocket.receive()
                    if msg.get("type") == "websocket.disconnect":
                        break
                    if msg.get("text") is not None:
                        await remote.send(msg["text"])
                    elif msg.get("bytes") is not None:
                        await remote.send(msg["bytes"])

            async def remote_to_client():
                async for msg in remote:
                    if isinstance(msg, bytes):
                        await websocket.send_bytes(msg)
                    else:
                        await websocket.send_text(msg)

            import asyncio
            tasks = [asyncio.create_task(client_to_remote()), asyncio.create_task(remote_to_client())]
            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for task in pending:
                task.cancel()
    except (WebSocketDisconnect, Exception):
        try:
            await websocket.close()
        except Exception:
            pass


@app.on_event("shutdown")
async def shutdown_client():
    await client.aclose()
