app.py - master - dsboikov/PyRelay

app.py
275 строк9 KB
Denis
Приложение переделано под POST-запросы
20 июл 2026, 23:44
20 июл 2026, 23:44170fba0
import ipaddress
import json
import os
import socket
import time
from collections import defaultdict, deque
from pathlib import Path
from urllib.parse import urlparse

import httpx
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

BASE_DIR = Path(__file__).resolve().parent
KEYS_FILE = BASE_DIR / "config" / "keys.json"
ALLOWLIST_FILE = BASE_DIR / "config" / "allowlist.json"

EMPTY_ALLOWLIST_MODE = os.environ.get("RELAY_EMPTY_ALLOWLIST_MODE", "allow")

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

ALLOWED_METHODS = {"GET", "POST", "PUT", "PATCH", "DELETE"}
BLOCKED_PORTS = {22, 25, 3306, 5432, 6379, 9200, 27017}

# content-encoding и content-length исключены отдельно: httpx сам разжимает
# тело ответа (gzip/br/deflate), поэтому заголовок Content-Encoding от
# апстрима после этого уже врёт — если пробросить его как есть, клиент
# попытается разжать уже разжатые данные и получит "incorrect header check".
# Мы всё равно оборачиваем тело в свой JSON, так что оригинальные заголовки
# ответа наружу не идут напрямую — только информационно, внутри envelope.
HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade", "host",
    "content-length", "content-encoding",
}


# ---------- Горячая перезагрузка конфигов ----------


class ReloadableConfig:
    """Перечитывает JSON-файл, если он изменился на диске. Без рестарта процесса."""

    def __init__(self, path: Path):
        self.path = path
        self._mtime = 0.0
        self._data = {}
        self._load()

    def _load(self):
        try:
            self._data = json.loads(self.path.read_text())
            self._mtime = self.path.stat().st_mtime
        except FileNotFoundError:
            self._data = {}
        except json.JSONDecodeError:
            pass  # оставляем последнюю валидную версию, не роняем сервис

    def get(self):
        try:
            mtime = self.path.stat().st_mtime
        except FileNotFoundError:
            return self._data
        if mtime != self._mtime:
            self._load()
        return self._data


keys_config = ReloadableConfig(KEYS_FILE)
allowlist_config = ReloadableConfig(ALLOWLIST_FILE)


def get_key_info(key: str):
    if not key:
        return None
    for entry in keys_config.get().get("keys", []):
        if entry.get("key") == key:
            return entry
    return None


def is_domain_allowed(hostname: str) -> bool:
    domains = allowlist_config.get().get("domains", [])
    if not domains:
        return EMPTY_ALLOWLIST_MODE == "allow"
    hostname = (hostname or "").lower()
    for d in domains:
        d = d.lower().strip()
        if not d:
            continue
        if hostname == d or hostname.endswith("." + d):
            return True
    return False


# ---------- Rate limiting ----------


def parse_rate(rate: str):
    count, _, period = rate.partition("/")
    count = int(count)
    seconds = {"second": 1, "minute": 60, "hour": 3600, "day": 86400}[period]
    return count, float(seconds)


class RateLimiter:
    def __init__(self):
        self._hits: dict[str, deque] = defaultdict(deque)

    def check(self, bucket_id: str, rate: str) -> bool:
        limit, window = parse_rate(rate)
        now = time.monotonic()
        q = self._hits[bucket_id]
        while q and now - q[0] > window:
            q.popleft()
        if len(q) >= limit:
            return False
        q.append(now)
        return True


limiter = RateLimiter()
IP_RATE_LIMIT = os.environ.get("RELAY_IP_RATE_LIMIT", "30/minute")


def get_client_ip(request: Request) -> str:
    return (
            request.headers.get("x-real-ip")
            or (request.client.host if request.client else "unknown")
    )


# ---------- SSRF-защита ----------


def is_private(host: str) -> bool:
    try:
        infos = socket.getaddrinfo(host, None)
    except socket.gaierror:
        return True  # не резолвится — считаем небезопасным
    for info in infos:
        ip = ipaddress.ip_address(info[4][0])
        if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved:
            return True
    return False


# ---------- Envelope-ответы ----------
# Каждый ответ relay — это JSON вида:
# {
#   "relay": { "stage": "...", "message": "..." },
#   "upstream": { "status": int, "headers": {...}, "body": ... } | null
# }
# Так всегда видно, на каком этапе конвейера произошла остановка: сам relay
# отказал (auth/rate-limit/allowlist/ssrf/validation) — тогда "upstream"
# будет null; либо запрос дошёл до цели — тогда в "upstream" лежит её
# реальный статус/заголовки/тело.


def relay_stopped(stage: str, message: str, http_status: int) -> JSONResponse:
    return JSONResponse(
        status_code=http_status,
        content={
            "relay": {"stage": stage, "message": message},
            "upstream": None,
        },
    )


# ---------- Основной эндпоинт ----------


@app.post("/relay")
async def relay(request: Request):
    client_ip = get_client_ip(request)

    try:
        payload = await request.json()
    except json.JSONDecodeError:
        return relay_stopped("invalid_request", "body is not valid JSON", 400)

    if not isinstance(payload, dict):
        return relay_stopped("invalid_request", "body must be a JSON object", 400)

    api_key = payload.get("x-relay-key") or request.headers.get("x-relay-key", "")
    key_info = get_key_info(api_key)
    if key_info is None:
        if not limiter.check(f"ip:{client_ip}", IP_RATE_LIMIT):
            return relay_stopped("rate_limited", "too many unauthorized attempts from this IP", 429)
        return relay_stopped("auth_failed", "missing or invalid x-relay-key", 401)

    rate = key_info.get("rate_limit", "60/minute")
    if not limiter.check(f"key:{api_key}", rate):
        return relay_stopped("rate_limited", f"rate limit exceeded for this key ({rate})", 429)

    target = payload.get("url")
    if not target or not isinstance(target, str):
        return relay_stopped("invalid_request", "field 'url' is required", 400)

    method = str(payload.get("type", "GET")).upper()
    if method not in ALLOWED_METHODS:
        return relay_stopped("invalid_request", f"unsupported method '{method}'", 400)

    parsed = urlparse(target)
    if parsed.scheme not in ("http", "https"):
        return relay_stopped("invalid_request", "url scheme must be http or https", 400)

    if not is_domain_allowed(parsed.hostname or ""):
        return relay_stopped("domain_blocked", f"target domain '{parsed.hostname}' is not in allowlist", 403)

    if is_private(parsed.hostname or ""):
        return relay_stopped("ssrf_blocked", "target host resolves to a private/internal address", 400)

    if parsed.port and parsed.port in BLOCKED_PORTS:
        return relay_stopped("ssrf_blocked", f"target port {parsed.port} is not allowed", 400)

    req_headers = payload.get("headers") or {}
    if not isinstance(req_headers, dict):
        return relay_stopped("invalid_request", "field 'headers' must be an object", 400)
    fwd_headers = {
        k: v for k, v in req_headers.items()
        if k.lower() not in HOP_BY_HOP and k.lower() != "x-relay-key"
    }

    req_body = payload.get("body", None)
    kwargs = {}
    if req_body is not None:
        if isinstance(req_body, (dict, list)):
            kwargs["json"] = req_body
        else:
            kwargs["content"] = str(req_body)

    async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as client:
        try:
            resp = await client.request(method, target, headers=fwd_headers, **kwargs)
        except httpx.RequestError as e:
            return relay_stopped("upstream_unreachable", f"{type(e).__name__}: {e}", 502)

    upstream_headers = {
        k: v for k, v in resp.headers.items() if k.lower() not in HOP_BY_HOP
    }

    # resp.content уже разжат httpx (если апстрим отдавал gzip/br/deflate) —
    # поэтому оригинальный content-encoding выше сознательно не пробрасывается.
    try:
        upstream_body = resp.json()
    except (json.JSONDecodeError, ValueError):
        try:
            upstream_body = resp.text
        except UnicodeDecodeError:
            upstream_body = None  # бинарный ответ — не пытаемся впихнуть в JSON

    return JSONResponse(
        status_code=200,  # сам relay отработал успешно; реальный статус цели — в upstream.status
        content={
            "relay": {"stage": "upstream_response", "message": "request completed"},
            "upstream": {
                "status": resp.status_code,
                "headers": upstream_headers,
                "body": upstream_body,
            },
        },
    )


@app.get("/health")
async def health():
    return {"status": "ok"}