rengine/shared/models/scan_engine.py

import uuid
from datetime import datetime
from functools import partial

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

from shared.definitions.constants import DEFAULT_GLOBAL_THREADS
from shared.models.scan_context import ScanContextCreate
from shared.utils.datetime import utc_now
from shared.utils.validation import clean_name, clean_optional_name


class ScanEngine(SQLModel, table=True):
    __tablename__ = "scan_engines"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True, index=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)
    intensity: str = Field(default="normal")
    global_threads: int = Field(default=DEFAULT_GLOBAL_THREADS)
    global_http_crawl: bool = Field(default=True)
    global_headers: list = Field(
        default_factory=list, sa_column=Column(JSON, nullable=False)
    )
    stages: dict = Field(default_factory=dict, sa_column=Column(JSON, nullable=False))
    yaml_source: str | None = Field(default=None, sa_column=Column(Text, nullable=True))
    tool_options: dict = Field(
        default_factory=dict, sa_column=Column(JSON, nullable=False)
    )
    created_at: datetime = Field(default_factory=utc_now)
    updated_at: datetime = Field(default_factory=utc_now)
    last_used_at: datetime | None = Field(default=None)


class ScanEngineCreate(BaseModel):
    name: str
    description: str | None = None
    intensity: str = "normal"
    global_threads: int = DEFAULT_GLOBAL_THREADS
    global_http_crawl: bool = True
    global_headers: list[str] = PydanticField(default_factory=list)
    stages: dict[str, dict] = PydanticField(default_factory=dict)
    yaml_source: str | None = None
    tool_options: dict[str, str] = PydanticField(default_factory=dict)

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


class ScanEngineUpdate(BaseModel):
    name: str | None = None
    description: str | None = None
    intensity: str | None = None
    global_threads: int | None = None
    global_http_crawl: bool | None = None
    global_headers: list[str] | None = None
    stages: dict[str, dict] | None = None
    yaml_source: str | None = None
    tool_options: dict[str, str] | None = None

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


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


class ScanEngineRead(BaseModel):
    id: uuid.UUID
    project_id: uuid.UUID
    created_by: uuid.UUID
    name: str
    description: str | None
    intensity: str
    global_threads: int
    global_http_crawl: bool
    global_headers: list[str]
    stages: dict[str, dict]
    yaml_source: str | None
    tool_options: dict[str, str] = PydanticField(default_factory=dict)
    usage: EngineUsage = PydanticField(default_factory=EngineUsage)
    created_at: datetime
    updated_at: datetime
    last_used_at: datetime | None


class StageField(BaseModel):
    name: str
    title: str
    description: str | None = None
    type: str
    default: object = None
    options: list[str] | None = None
    option_labels: dict[str, str] | None = None
    minimum: int | None = None
    maximum: int | None = None
    scale: str | None = None
    widget: str | None = None
    # what a widget selects, when the widget needs it (wordlist kind)
    kind: str | None = None
    launch: bool = False


class StageCatalogEntry(BaseModel):
    name: str
    title: str
    description: str
    phase: str
    level: int
    applies_to: list[str]
    tools: list[str]
    api_keys: list[str]
    requires_api_keys: bool = False
    touches_target: bool = True
    always_on: bool = False
    launch_fields: list[str] = PydanticField(default_factory=list)
    group: str
    role: str
    consumes: list[str] = PydanticField(default_factory=list)
    produces: list[str] = PydanticField(default_factory=list)
    defaults: dict
    fields: list[StageField] = PydanticField(default_factory=list)


class StageGroupEntry(BaseModel):
    key: str
    label: str


class ToolOption(BaseModel):
    name: str
    label: str
    phase: str
    example: str


class EnginePreset(BaseModel):
    name: str
    title: str
    description: str
    stages: dict[str, dict]


class EngineCatalog(BaseModel):
    phases: list[str]
    stages: list[StageCatalogEntry]
    rate_tools: list[str]
    tool_options: list[ToolOption]
    presets: list[EnginePreset]
    target_types: list[str]
    groups: list[StageGroupEntry] = PydanticField(default_factory=list)
    seed_produces: dict[str, list[str]] = PydanticField(default_factory=dict)


class PreviewResolved(BaseModel):
    header_names: list[str] = PydanticField(default_factory=list)
    global_threads: int = DEFAULT_GLOBAL_THREADS
    global_rate_limit_ceiling: int | None = None
    per_tool_rate_limits: dict[str, int] = PydanticField(default_factory=dict)
    excluded_subdomains: list[str] = PydanticField(default_factory=list)
    excluded_paths: list[str] = PydanticField(default_factory=list)
    excluded_ips: list[str] = PydanticField(default_factory=list)
    included_subdomains: list[str] = PydanticField(default_factory=list)
    follow_redirects: bool | None = None
    http_protocol: str = "both"


class EnginePreviewResult(BaseModel):
    phases: list = PydanticField(default_factory=list)
    resolved_stages: dict[str, dict] = PydanticField(default_factory=dict)
    resolved: PreviewResolved = PydanticField(default_factory=PreviewResolved)
    warnings: list[str] = PydanticField(default_factory=list)


class EnginePreviewRequest(BaseModel):
    target_type: str
    context_id: uuid.UUID | None = None
    context: ScanContextCreate | None = None
    intensity: str = "normal"
    global_threads: int = DEFAULT_GLOBAL_THREADS
    stages: dict[str, dict] = PydanticField(default_factory=dict)