rengine/shared/models/scan_context.py

import uuid
from datetime import datetime
from functools import partial

from pydantic import BaseModel, field_validator
from sqlalchemy import Column
from sqlalchemy.types import JSON
from sqlmodel import Field, SQLModel

from shared.models.types import EncryptedJSON
from shared.utils.datetime import utc_now
from shared.utils.validation import clean_name, clean_optional_name

AUTH_TYPES = ("none", "header", "bearer", "basic", "cookie", "api_key")
HTTP_PROTOCOLS = ("both", "http_only", "https_only")

# the scheme a probe is restricted to; "both" leaves the choice to the tool
PROBE_SCHEME: dict[str, str] = {"http_only": "http", "https_only": "https"}


def valid_rate_tools() -> tuple[str, ...]:
    from stages.registry import rate_tools  # noqa: PLC0415

    return rate_tools()


MULTIPLIERS = (0.5, 1.0, 2.0)


class AuthConfig(BaseModel):
    auth_type: str = "none"
    bearer_token: str | None = None
    basic_username: str | None = None
    basic_password: str | None = None
    header_name: str | None = None
    header_value: str | None = None
    cookie_value: str | None = None
    api_key_name: str | None = None
    api_key_value: str | None = None


class AuthHeader(BaseModel):
    name: str
    value: str


class ScanContext(SQLModel, table=True):
    __tablename__ = "scan_contexts"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    project_id: uuid.UUID = Field(foreign_key="projects.id", index=True)
    created_by: uuid.UUID = Field(foreign_key="users.id")
    name: str = Field(max_length=200)
    description: str | None = Field(default=None, max_length=1000)
    auth_type: str = Field(default="none")
    auth: dict = Field(
        default_factory=dict, sa_column=Column(EncryptedJSON, nullable=False)
    )
    extra_headers: list = Field(
        default_factory=list, sa_column=Column(EncryptedJSON, nullable=False)
    )
    global_rate_limit_override: int | None = Field(default=None)
    per_tool_rate_overrides: dict = Field(
        default_factory=dict, sa_column=Column(JSON, nullable=False)
    )
    thread_multiplier: float = Field(default=1.0)
    timeout_multiplier: float = Field(default=1.0)
    excluded_subdomains: list = Field(
        default_factory=list, sa_column=Column(JSON, nullable=False)
    )
    excluded_paths: list = Field(
        default_factory=list, sa_column=Column(JSON, nullable=False)
    )
    excluded_ips: list = Field(
        default_factory=list, sa_column=Column(JSON, nullable=False)
    )
    included_subdomains: list = Field(
        default_factory=list, sa_column=Column(JSON, nullable=False)
    )
    follow_redirects_override: bool | None = Field(default=None)
    http_protocol: str = Field(default="both")
    proxy_id: uuid.UUID | None = Field(default=None, index=True)
    created_at: datetime = Field(default_factory=utc_now)
    updated_at: datetime = Field(default_factory=utc_now)
    last_used_at: datetime | None = Field(default=None)
    last_used_scan_id: uuid.UUID | None = Field(default=None)


class ScanContextCreate(BaseModel):
    name: str
    description: str | None = None
    auth_type: str = "none"
    auth: AuthConfig | None = None
    extra_headers: list[AuthHeader] = Field(default_factory=list)
    global_rate_limit_override: int | None = None
    per_tool_rate_overrides: dict[str, int] = Field(default_factory=dict)
    thread_multiplier: float = 1.0
    timeout_multiplier: float = 1.0
    excluded_subdomains: list[str] = Field(default_factory=list)
    excluded_paths: list[str] = Field(default_factory=list)
    excluded_ips: list[str] = Field(default_factory=list)
    included_subdomains: list[str] = Field(default_factory=list)
    follow_redirects_override: bool | None = None
    http_protocol: str = "both"
    proxy_id: uuid.UUID | None = None

    _validate_name = field_validator("name")(partial(clean_name, max_len=200))


class ScanContextUpdate(BaseModel):
    name: str | None = None
    description: str | None = None
    auth_type: str | None = None
    auth: AuthConfig | None = None
    extra_headers: list[AuthHeader] | None = None
    global_rate_limit_override: int | None = None
    per_tool_rate_overrides: dict[str, int] | None = None
    thread_multiplier: float | None = None
    timeout_multiplier: float | None = None
    excluded_subdomains: list[str] | None = None
    excluded_paths: list[str] | None = None
    excluded_ips: list[str] | None = None
    included_subdomains: list[str] | None = None
    follow_redirects_override: bool | None = None
    http_protocol: str | None = None
    proxy_id: uuid.UUID | None = None

    _validate_name = field_validator("name")(partial(clean_optional_name, max_len=200))


class ContextUsage(BaseModel):
    schedules: int = 0
    scans: int = 0


class ScanContextRead(BaseModel):
    id: uuid.UUID
    project_id: uuid.UUID
    created_by: uuid.UUID
    name: str
    description: str | None
    auth_type: str
    auth: AuthConfig
    auth_summary: str
    extra_headers: list[AuthHeader]
    global_rate_limit_override: int | None
    per_tool_rate_overrides: dict[str, int]
    thread_multiplier: float
    timeout_multiplier: float
    excluded_subdomains: list[str]
    excluded_paths: list[str]
    excluded_ips: list[str]
    included_subdomains: list[str]
    follow_redirects_override: bool | None
    http_protocol: str
    proxy_id: uuid.UUID | None
    created_at: datetime
    updated_at: datetime
    last_used_at: datetime | None
    last_used_scan_id: uuid.UUID | None
    usage: ContextUsage = Field(default_factory=ContextUsage)