mi
273 lines
10 KiB
Python
273 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import hmac
|
|
import json
|
|
import uuid
|
|
from time import perf_counter
|
|
from typing import Annotated
|
|
|
|
from fastapi import Depends, FastAPI, Header, Request
|
|
from fastapi.responses import JSONResponse
|
|
from pydantic import TypeAdapter, ValidationError
|
|
|
|
from app.contracts import CheckRequest, ErrorBody, ErrorEnvelope, Pending, Verdict
|
|
from app.db import TaskStatus
|
|
from app.rate_limit import RateLimited
|
|
from app.repository import ConflictError
|
|
from app.service import CapabilityUnavailable, SafetyService, TaskFailed
|
|
from app.telemetry import logger, record_poll
|
|
|
|
CHECK_ADAPTER = TypeAdapter(CheckRequest)
|
|
MAX_BODY = 16_384
|
|
|
|
|
|
def error(
|
|
status: int, code: str, request_id: str, details: dict[str, object] | None = None
|
|
) -> JSONResponse:
|
|
body = ErrorEnvelope(
|
|
error=ErrorBody(
|
|
code=code,
|
|
message={
|
|
"validation_error": "Request is invalid",
|
|
"service_unauthorized": "Service authentication failed",
|
|
}.get(code, "Request could not be completed"),
|
|
request_id=request_id,
|
|
details=details or {},
|
|
)
|
|
)
|
|
return JSONResponse(status_code=status, content=body.model_dump(mode="json"))
|
|
|
|
|
|
def create_app(service: SafetyService, token: str) -> FastAPI:
|
|
app = FastAPI(title="HAN Message Safety", version="2.0.0", docs_url=None, redoc_url=None)
|
|
|
|
async def authenticate(
|
|
request: Request,
|
|
provided: Annotated[str | None, Header(alias="X-Service-Token")] = None,
|
|
) -> None:
|
|
if not provided or not hmac.compare_digest(provided.encode(), token.encode()):
|
|
request.state.auth_failed = True
|
|
raise PermissionError
|
|
|
|
@app.exception_handler(PermissionError)
|
|
async def auth_error(request: Request, _: PermissionError) -> JSONResponse:
|
|
return error(401, "service_unauthorized", request.state.request_id)
|
|
|
|
@app.middleware("http")
|
|
async def request_context(request: Request, call_next):
|
|
started = perf_counter()
|
|
supplied = request.headers.get("X-Request-ID")
|
|
try:
|
|
request.state.request_id = str(uuid.UUID(supplied)) if supplied else str(uuid.uuid4())
|
|
except ValueError:
|
|
request.state.request_id = str(uuid.uuid4())
|
|
try:
|
|
response = await call_next(request)
|
|
except Exception:
|
|
logger.error(
|
|
"Safety request failed",
|
|
extra={
|
|
"event": "http_request_failed",
|
|
"request_id": request.state.request_id,
|
|
"route": _route_template(request),
|
|
"method": request.method,
|
|
},
|
|
)
|
|
raise
|
|
response.headers["X-Request-ID"] = request.state.request_id
|
|
response.headers["Cache-Control"] = "no-store"
|
|
if request.url.path != "/health/live":
|
|
logger.info(
|
|
"Safety request completed",
|
|
extra={
|
|
"event": "http_request_completed",
|
|
"request_id": request.state.request_id,
|
|
"route": _route_template(request),
|
|
"method": request.method,
|
|
"status_code": response.status_code,
|
|
"duration_ms": round((perf_counter() - started) * 1000, 3),
|
|
},
|
|
)
|
|
return response
|
|
|
|
@app.get("/health/live")
|
|
async def live() -> dict[str, str]:
|
|
return {"status": "ok"}
|
|
|
|
@app.get("/health/ready")
|
|
async def ready() -> JSONResponse:
|
|
await service.refresh_antivirus()
|
|
mode = "mock" if service.mode.mock else "standard"
|
|
components = {
|
|
"postgres": "ok",
|
|
"redis": "degraded",
|
|
"s3_quarantine": "bypassed"
|
|
if service.mode.mock
|
|
else ("ok" if service.files_ready else "down"),
|
|
"worker": "bypassed" if service.mode.mock else "ok",
|
|
"antivirus": "bypassed"
|
|
if service.mode.mock
|
|
else ("ok" if service.files_ready else "down"),
|
|
"dns": "bypassed" if service.mode.mock else ("ok" if service.links_ready else "down"),
|
|
"rules": "bypassed" if service.mode.mock else "ok",
|
|
}
|
|
capabilities = {
|
|
"text": "ready",
|
|
"links": "bypassed"
|
|
if service.mode.mock
|
|
else ("ready" if service.links_ready else "unavailable"),
|
|
"files": "bypassed"
|
|
if service.mode.mock
|
|
else ("ready" if service.files_ready else "unavailable"),
|
|
"worker": "bypassed" if service.mode.mock else "ready",
|
|
}
|
|
body: dict[str, object] = {
|
|
"status": "degraded"
|
|
if service.mode.mock or "degraded" in components.values()
|
|
else "ok",
|
|
"processing_mode": mode,
|
|
"config_version": service.config.version,
|
|
"components": components,
|
|
"capabilities": capabilities,
|
|
}
|
|
if service.mode.mock:
|
|
body["mock_policy"] = {
|
|
"text": "allow" if service.mode.text_free else "deny",
|
|
"file": "allow" if service.mode.file_free else "deny",
|
|
}
|
|
return JSONResponse(content=body)
|
|
|
|
@app.post("/internal/safety/v2/messages/check", dependencies=[Depends(authenticate)])
|
|
async def check(request: Request) -> JSONResponse:
|
|
content_type = request.headers.get("content-type", "").lower().replace(" ", "")
|
|
if content_type not in {"application/json", "application/json;charset=utf-8"}:
|
|
return error(
|
|
400,
|
|
"validation_error",
|
|
request.state.request_id,
|
|
{"field": "content-type", "constraint": "application/json; charset=utf-8"},
|
|
)
|
|
body = await request.body()
|
|
if len(body) > MAX_BODY:
|
|
return error(
|
|
400,
|
|
"validation_error",
|
|
request.state.request_id,
|
|
{"field": "body", "constraint": "max_bytes"},
|
|
)
|
|
try:
|
|
payload = CHECK_ADAPTER.validate_json(body, strict=True)
|
|
result = await service.check(payload)
|
|
except (ValidationError, json.JSONDecodeError, ValueError):
|
|
return error(
|
|
400,
|
|
"validation_error",
|
|
request.state.request_id,
|
|
{"field": "body", "constraint": "strict_dto"},
|
|
)
|
|
except ConflictError:
|
|
message_id = "unknown"
|
|
try:
|
|
message_id = str(json.loads(body).get("message_id", "unknown"))
|
|
except (ValueError, AttributeError):
|
|
pass
|
|
return error(
|
|
409,
|
|
"safety_request_conflict",
|
|
request.state.request_id,
|
|
{"message_id": message_id, "terminal": True, "retryable": False},
|
|
)
|
|
except CapabilityUnavailable as exc:
|
|
return error(
|
|
503,
|
|
"dependency_unavailable",
|
|
request.state.request_id,
|
|
{"dependency_category": exc.category, "terminal": False, "retryable": True},
|
|
)
|
|
except RateLimited as exc:
|
|
response = error(
|
|
429,
|
|
"rate_limit_exceeded",
|
|
request.state.request_id,
|
|
{"retryable": True, "retry_after_sec": exc.retry_after},
|
|
)
|
|
response.headers["Retry-After"] = str(exc.retry_after)
|
|
return response
|
|
except TaskFailed as exc:
|
|
return error(
|
|
503,
|
|
"task_failed",
|
|
request.state.request_id,
|
|
{
|
|
"task_id": str(exc.task_id),
|
|
"task_status": "failed",
|
|
"terminal": True,
|
|
"retryable": False,
|
|
},
|
|
)
|
|
status = 202 if isinstance(result, Pending) else (200 if result.verdict == "allow" else 403)
|
|
response = JSONResponse(status_code=status, content=result.model_dump(mode="json"))
|
|
if isinstance(result, Pending):
|
|
response.headers["Location"] = f"/internal/safety/v2/messages/tasks/{result.task_id}"
|
|
response.headers["Retry-After"] = str(result.poll_after_ms // 1000)
|
|
return response
|
|
|
|
@app.get("/internal/safety/v2/messages/tasks/{task_id}", dependencies=[Depends(authenticate)])
|
|
async def get_task(request: Request, task_id: str) -> JSONResponse:
|
|
try:
|
|
parsed = uuid.UUID(task_id)
|
|
except ValueError:
|
|
return error(
|
|
400,
|
|
"validation_error",
|
|
request.state.request_id,
|
|
{"field": "task_id", "constraint": "uuid"},
|
|
)
|
|
task = await service.repository.task(parsed)
|
|
if not task:
|
|
record_poll("not_found")
|
|
return error(404, "task_not_found", request.state.request_id)
|
|
if task.status == TaskStatus.failed:
|
|
record_poll("failed")
|
|
return error(
|
|
503,
|
|
"task_failed",
|
|
request.state.request_id,
|
|
{
|
|
"task_id": str(task.id),
|
|
"task_status": "failed",
|
|
"terminal": True,
|
|
"retryable": False,
|
|
},
|
|
)
|
|
if task.status in {TaskStatus.pending, TaskStatus.processing}:
|
|
record_poll(task.status.value)
|
|
result: Pending | Verdict = Pending(
|
|
config_version=task.config_version,
|
|
task_id=task.id,
|
|
expires_at=task.expires_at,
|
|
rules_version=task.rules_version,
|
|
)
|
|
else:
|
|
record_poll("terminal")
|
|
result = service._verdict(
|
|
task.status == TaskStatus.allowed,
|
|
task.processing_mode,
|
|
task.rule_id or "safety.all_checks_passed",
|
|
task.rules_version,
|
|
config_version=task.config_version,
|
|
)
|
|
status = 202 if isinstance(result, Pending) else (200 if result.verdict == "allow" else 403)
|
|
response = JSONResponse(status_code=status, content=result.model_dump(mode="json"))
|
|
if isinstance(result, Pending):
|
|
response.headers["Location"] = str(request.url.path)
|
|
response.headers["Retry-After"] = "2"
|
|
return response
|
|
|
|
return app
|
|
|
|
|
|
def _route_template(request: Request) -> str:
|
|
route = request.scope.get("route")
|
|
return getattr(route, "path", "unmatched")
|