rengine/api/app/services/scan_engine/validation.py

import shlex

import yaml
from fastapi import HTTPException, status
from pydantic import ValidationError

from shared.definitions.tools import MAX_TOOL_OPTION_LEN, TOOL_NAMES
from shared.enums.scan import INTENSITIES
from shared.services.scan_resolve import (
    _SENSITIVE_HEADER,
    MASK,
    _reject_ctrl,
    redact_command,
    secret_runs,
)
from shared.utils.yaml_safe import DocumentTooLargeError, load_document
from stages.registry import stage_by_name, stages

_MAX_HEADERS = 1000
_MAX_HEADER_LEN = 4096
_MAX_YAML_LEN = 512 * 1024
_INTENSITIES = set(INTENSITIES)
_MAX_ENGINE_THREADS = 1000


def _mask_tool_options(options: dict | None) -> dict[str, str]:
    return {t: redact_command(v) for t, v in (options or {}).items()}


def _unmask_tool_options(submitted: dict | None, stored: dict | None) -> dict[str, str]:
    """Restore each masked run in place."""
    stored = stored or {}
    out: dict[str, str] = {}
    for tool, value in (submitted or {}).items():
        if not value or MASK not in value or tool not in stored:
            out[tool] = value
            continue
        originals = secret_runs(stored[tool])
        restored = value
        for original in originals:
            restored = restored.replace(MASK, original, 1)
        out[tool] = stored[tool] if MASK in restored else restored
    return out


_MAX_HEADERS = 1000
_MAX_HEADER_LEN = 4096
_MAX_YAML_LEN = 512 * 1024


def _validate_yaml_source(source: str | None) -> str | None:
    if source is None:
        return None
    if len(source) > _MAX_YAML_LEN:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Stage YAML may not exceed {_MAX_YAML_LEN} bytes.",
        )
    try:
        load_document(source)
    except DocumentTooLargeError as exc:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)
        ) from exc
    except yaml.YAMLError as exc:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid YAML: {exc}"
        ) from exc
    return source


_INTENSITIES = set(INTENSITIES)
_MAX_ENGINE_THREADS = 1000


def _validate_tool_options(options: dict | None) -> dict[str, str]:
    """Keep only known tools."""
    clean: dict[str, str] = {}
    for tool, raw in (options or {}).items():
        if tool not in TOOL_NAMES:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=(
                    f"'{tool}' does not take custom arguments. "
                    f"Tools that do: {', '.join(sorted(TOOL_NAMES))}."
                ),
            )
        value = (raw or "").strip()
        if not value:
            continue
        if len(value) > MAX_TOOL_OPTION_LEN:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"{tool} options may not exceed {MAX_TOOL_OPTION_LEN} characters.",
            )
        _reject_ctrl(f"{tool} options", value)
        try:
            shlex.split(value)
        except ValueError as exc:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"{tool} options are not valid shell arguments: {exc}",
            ) from exc
        clean[tool] = value
    return clean


def _validate_intensity(intensity: str | None) -> None:
    if intensity is not None and intensity not in _INTENSITIES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"intensity must be one of {sorted(_INTENSITIES)}.",
        )


def _validate_global_threads(threads: int | None) -> None:
    if threads is not None and not (1 <= threads <= _MAX_ENGINE_THREADS):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"global_threads must be between 1 and {_MAX_ENGINE_THREADS}.",
        )


def _validate_global_headers(headers: list) -> None:
    if headers and len(headers) > _MAX_HEADERS:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"global_headers may not exceed {_MAX_HEADERS} entries.",
        )
    for line in headers or []:
        if isinstance(line, str) and len(line) > _MAX_HEADER_LEN:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Each global header may not exceed {_MAX_HEADER_LEN} characters.",
            )
        if not isinstance(line, str) or ":" not in line:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Each global header must be a 'Name: Value' string.",
            )
        name, value = line.split(":", 1)
        if not name.strip():
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Each global header must have a non-empty name.",
            )
        _reject_ctrl("Header name", name)
        _reject_ctrl("Header value", value)


def _mask_global_headers(headers: list) -> list[str]:
    out: list[str] = []
    for line in headers or []:
        if isinstance(line, str) and ":" in line:
            name, value = line.split(":", 1)
            if value.strip() and _SENSITIVE_HEADER.search(name.strip()):
                out.append(f"{name.strip()}: {MASK}")
                continue
        out.append(line)
    return out


def _unmask_global_headers(incoming: list, stored: list) -> list[str]:
    stored_map: dict[str, str] = {}
    for line in stored or []:
        if isinstance(line, str) and ":" in line:
            name, value = line.split(":", 1)
            stored_map[name.strip().lower()] = value.strip()
    out: list[str] = []
    for line in incoming or []:
        if isinstance(line, str) and ":" in line:
            name, value = line.split(":", 1)
            if value.strip() == MASK and name.strip().lower() in stored_map:
                out.append(f"{name.strip()}: {stored_map[name.strip().lower()]}")
                continue
        out.append(line)
    return out


def _validate_stages(submitted: dict | None) -> dict[str, dict]:
    known = stage_by_name()
    clean: dict[str, dict] = {}
    for name, authored in (submitted or {}).items():
        spec = known.get(name)
        raw = authored
        if spec is not None and spec.catalog_hidden:
            continue
        if spec is not None and spec.always_on and isinstance(raw, dict):
            raw = {k: v for k, v in raw.items() if k != "enabled"}
            if not raw:
                continue
        if spec is None:
            offered = sorted(n for n, s in known.items() if not s.catalog_hidden)
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Unknown stage '{name}'. Known stages: {', '.join(offered)}.",
            )
        if not isinstance(raw, dict):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Stage '{name}' config must be an object.",
            )
        known_fields = set(spec.config_model.model_fields)
        unknown = [k for k in raw if k not in known_fields]
        if unknown:
            hint = ", ".join(sorted(known_fields))
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=(
                    f"Stage '{name}' has no setting {', '.join(repr(u) for u in unknown)}. "
                    f"Valid settings: {hint}."
                ),
            )
        try:
            validated = spec.config_model(**raw).model_dump()
            clean[name] = {k: validated[k] for k in raw if k in validated}
        except ValidationError as exc:
            problems = "; ".join(
                f"{'.'.join(str(p) for p in e['loc'])}: {e['msg']}"
                for e in exc.errors()
            )
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Invalid config for stage '{name}': {problems}",
            ) from exc
    return clean


def _full_stages(stored: dict | None) -> dict[str, dict]:
    stored = stored or {}
    return {
        spec.name: spec.config_model(**(stored.get(spec.name) or {})).model_dump(
            mode="json"
        )
        for spec in stages()
        if not spec.catalog_hidden
    }


_JSON_TYPES = {
    "integer": "integer",
    "number": "number",
    "boolean": "boolean",
    "string": "string",
}