rengine/shared/models/scan.py

import uuid
from datetime import datetime

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

from shared.definitions.constants import MAX_SCAN_BATCH
from shared.definitions.rescan import MAX_RUN_ASSETS, MAX_RUN_SCANS, SeedKind
from shared.definitions.surface import SURFACE_ORDER
from shared.enums.scan import ScanScope, ScanStatus
from shared.services.scan_resolve import ResolvedScanConfig
from shared.utils.datetime import utc_now

SCAN_STATUSES = tuple(s.value for s in ScanStatus)


class SeedAsset(BaseModel):
    """One asset a focused scan starts from."""

    model_config = ConfigDict(extra="forbid")

    kind: str = Field(max_length=16)
    value: str = Field(max_length=500)

    @model_validator(mode="after")
    def _known_kind(self):
        if self.kind not in {k.value for k in SeedKind}:
            msg = f"Unknown seed kind '{self.kind}'."
            raise ValueError(msg)
        if not self.value.strip():
            msg = "Seed asset value is required."
            raise ValueError(msg)
        return self


class SeedPick(BaseModel):
    model_config = ConfigDict(extra="forbid")

    value: str = Field(max_length=500)
    scan_id: uuid.UUID | None = None


class QuerySelection(BaseModel):
    model_config = ConfigDict(extra="forbid")

    filter: dict = Field(default_factory=dict)
    scan_ids: list[uuid.UUID] = Field(default_factory=list, max_length=MAX_RUN_SCANS)


class SeedSelection(BaseModel):
    model_config = ConfigDict(extra="forbid")

    dimension: str = Field(max_length=32)
    picks: list[SeedPick] = Field(default_factory=list, max_length=MAX_RUN_ASSETS)
    query: QuerySelection | None = None
    exclude: list[str] = Field(default_factory=list, max_length=MAX_RUN_ASSETS)

    @model_validator(mode="after")
    def _one_form(self):
        if self.dimension not in SURFACE_ORDER:
            msg = f"Unknown dimension '{self.dimension}'."
            raise ValueError(msg)
        if bool(self.picks) == (self.query is not None):
            msg = "Provide either picks or a query, not both."
            raise ValueError(msg)
        return self


class SeedGroup(BaseModel):
    """Seed values for one target's scan."""

    scan_id: uuid.UUID
    target_id: uuid.UUID
    target_value: str
    values: list[str] = Field(default_factory=list)


class ResolvedSelection(BaseModel):
    groups: list[SeedGroup] = Field(default_factory=list)
    total: int = 0
    capped: bool = False
    matched: int | None = None


class SeedGroupSummary(BaseModel):
    target_id: uuid.UUID
    target_value: str
    scan_id: uuid.UUID
    count: int


class RunPreview(BaseModel):
    dimension: str
    seed_kind: str
    asset_count: int = 0
    matched: int | None = None
    target_count: int = 0
    capped: bool = False
    targets: list[SeedGroupSummary] = Field(default_factory=list)
    stage_titles: list[str] = Field(default_factory=list)


class FocusedRunRead(BaseModel):
    """The scans one rescan created."""

    run_group_id: uuid.UUID
    scans: list["ScanRead"] = Field(default_factory=list)
    asset_count: int = 0
    matched: int | None = None
    target_count: int = 0
    capped: bool = False
    stage_titles: list[str] = Field(default_factory=list)


class Scan(SQLModel, table=True):
    __tablename__ = "scans"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    project_id: uuid.UUID = Field(foreign_key="projects.id", index=True)
    target_id: uuid.UUID = Field(
        foreign_key="targets.id", index=True, ondelete="CASCADE"
    )
    engine_id: uuid.UUID | None = Field(default=None, index=True)
    engine_name: str = Field(max_length=200)
    context_id: uuid.UUID | None = Field(default=None)
    context_name: str | None = Field(default=None, max_length=200)
    schedule_id: uuid.UUID | None = Field(default=None, index=True)
    schedule_type: str | None = Field(default=None, max_length=20)
    scope: str = Field(default=ScanScope.FULL.value, max_length=16, index=True)
    parent_scan_id: uuid.UUID | None = Field(default=None, index=True)
    run_group_id: uuid.UUID | None = Field(default=None, index=True)
    execution_config: dict = Field(sa_column=Column(JSON, nullable=False))
    status: str = Field(default=ScanStatus.PENDING.value, index=True)
    celery_task_ids: list = Field(
        default_factory=list, sa_column=Column(JSON, nullable=False)
    )
    run_epoch: int = Field(default=0)
    paused_at: datetime | None = Field(default=None)
    paused_seconds: float = Field(default=0.0)
    subdomains_found: int = Field(default=0)
    ips_found: int = Field(default=0)
    open_ports_found: int = Field(default=0)
    http_assets_found: int = Field(default=0)
    vulnerabilities_found: int = Field(default=0)
    endpoints_found: int = Field(default=0)
    error: str | None = Field(default=None, max_length=2000)
    interest_signature: str | None = Field(default=None, max_length=64)
    interest_judged_at: datetime | None = Field(default=None)
    interest_model: str | None = Field(default=None, max_length=80)
    created_by: uuid.UUID = Field(foreign_key="users.id")
    created_at: datetime = Field(default_factory=utc_now)
    started_at: datetime | None = Field(default=None)
    completed_at: datetime | None = Field(default=None)


def fold_pause(scan: Scan, at: datetime) -> None:
    """Move an open pause into paused_seconds."""
    if scan.paused_at is None:
        return
    scan.paused_seconds = (scan.paused_seconds or 0.0) + (
        at - scan.paused_at
    ).total_seconds()
    scan.paused_at = None


class ScanCreate(BaseModel):
    """A launch names a saved engine, or runs an ad hoc plan carried in `overrides`."""

    model_config = ConfigDict(extra="forbid")

    engine_id: uuid.UUID | None = None
    context_id: uuid.UUID | None = None
    target_id: uuid.UUID | None = None
    target_value: str | None = Field(default=None, max_length=500)
    overrides: dict[str, dict] = Field(default_factory=dict)
    intensity: str | None = None
    seed_assets: list[SeedAsset] = Field(
        default_factory=list, max_length=MAX_RUN_ASSETS
    )
    parent_scan_id: uuid.UUID | None = None
    run_group_id: uuid.UUID | None = None
    dimension: str | None = Field(default=None, max_length=32)

    @model_validator(mode="after")
    def _require_target(self):
        if (self.target_id is None) == (not self.target_value):
            msg = "Provide either target_id or target_value."
            raise ValueError(msg)
        return self


class RescanCreate(BaseModel):
    """Re-run chosen stages against assets picked from an earlier run."""

    model_config = ConfigDict(extra="forbid")

    parent_scan_id: uuid.UUID | None = None
    dimension: str = Field(default="", max_length=32)
    assets: list[str] = Field(default_factory=list, max_length=MAX_RUN_ASSETS)
    selection: SeedSelection | None = None
    stages: list[str] = Field(default_factory=list, max_length=20)
    overrides: dict[str, dict] = Field(default_factory=dict)
    context_id: uuid.UUID | None = None
    intensity: str | None = None
    template_ids: list[str] = Field(default_factory=list, max_length=50)

    @model_validator(mode="after")
    def _one_selection(self):
        """Legacy parent_scan_id plus assets normalizes to selection."""
        legacy = self.parent_scan_id is not None or bool(self.assets)
        if legacy == (self.selection is not None):
            msg = "Provide either a selection, or a parent scan and its assets."
            raise ValueError(msg)
        if self.selection is not None:
            self.dimension = self.selection.dimension
            return self
        if self.dimension not in SURFACE_ORDER:
            msg = f"Unknown dimension '{self.dimension}'."
            raise ValueError(msg)
        if self.parent_scan_id is None:
            msg = "Provide the scan these assets came from."
            raise ValueError(msg)
        cleaned = list(dict.fromkeys(a.strip() for a in self.assets if a and a.strip()))
        if not cleaned:
            msg = "Select at least one asset to rescan."
            raise ValueError(msg)
        self.selection = SeedSelection(
            dimension=self.dimension,
            picks=[
                SeedPick(value=value, scan_id=self.parent_scan_id) for value in cleaned
            ],
        )
        self.assets = cleaned
        return self


class RescanDimension(BaseModel):
    dimension: str
    label: str
    noun: str
    noun_plural: str
    seed_kind: str
    default_stages: list[str]


class RescanSchema(BaseModel):
    dimensions: list[RescanDimension]
    rescannable_stages: list[str]
    max_assets: int
    max_scans: int


class ScanBatchCreate(BaseModel):
    model_config = ConfigDict(extra="forbid")

    engine_id: uuid.UUID | None = None
    context_id: uuid.UUID | None = None
    target_ids: list[uuid.UUID] = Field(default_factory=list, max_length=MAX_SCAN_BATCH)
    target_values: list[str] = Field(default_factory=list, max_length=MAX_SCAN_BATCH)
    overrides: dict[str, dict] = Field(default_factory=dict)
    intensity: str | None = None

    @model_validator(mode="after")
    def _require_targets(self):
        total = len(self.target_ids) + len(self.target_values)
        if total == 0:
            msg = "Select at least one target."
            raise ValueError(msg)
        if total > MAX_SCAN_BATCH:
            msg = f"Select at most {MAX_SCAN_BATCH} targets."
            raise ValueError(msg)
        return self


class ScanRead(BaseModel):
    id: uuid.UUID
    project_id: uuid.UUID
    target_id: uuid.UUID
    engine_id: uuid.UUID | None
    engine_name: str
    context_id: uuid.UUID | None
    context_name: str | None
    schedule_id: uuid.UUID | None = None
    schedule_type: str | None = None
    scope: str = ScanScope.FULL.value
    parent_scan_id: uuid.UUID | None = None
    run_group_id: uuid.UUID | None = None
    seed_count: int = 0
    execution_config: ResolvedScanConfig
    auth_summary: str
    status: str
    subdomains_found: int
    ips_found: int
    open_ports_found: int
    http_assets_found: int
    vulnerabilities_found: int
    endpoints_found: int
    error: str | None
    created_by: uuid.UUID
    created_at: datetime
    started_at: datetime | None
    completed_at: datetime | None
    paused_at: datetime | None = None
    paused_seconds: float = 0.0
    duration_seconds: float | None = None
    new_subdomains: int | None = None
    gone_subdomains: int | None = None
    prev_subdomains_found: int | None = None
    is_first_scan: bool | None = None


class ScanStatusCounts(BaseModel):
    pending: int = 0
    running: int = 0
    paused: int = 0
    completed: int = 0
    failed: int = 0
    cancelled: int = 0


class ScanDailyCount(BaseModel):
    date: str
    count: int
    completed: int = 0
    failed: int = 0
    cancelled: int = 0
    running: int = 0
    paused: int = 0
    pending: int = 0
    new_subdomains: int = 0


class ScanFacet(BaseModel):
    name: str
    count: int


class ScanChanges(BaseModel):
    window: str
    new_subdomains: int
    retired_subdomains: int = 0
    targets_changed: int
    scans_run: int
    failed_runs: int


class ScanTargetGroup(BaseModel):
    target_id: uuid.UUID
    target_value: str
    target_type: str
    scan_count: int
    last_scan_at: datetime
    last_status: str
    running: int
    trend: list[int] = []


class ScanExportRow(BaseModel):
    target: str
    status: str
    engine: str
    context: str | None
    schedule_type: str | None = None
    subdomains: int
    ips: int
    open_ports: int
    vulnerabilities: int
    endpoints: int
    duration_seconds: float | None
    started_at: datetime | None
    completed_at: datetime | None
    created_at: datetime


class ScanStats(BaseModel):
    total: int
    running: int
    by_status: ScanStatusCounts
    last_scan_at: datetime | None
    avg_duration_seconds: float | None
    success_rate: float | None
    daily: list[ScanDailyCount]
    engines: list[ScanFacet] = []
    contexts: list[ScanFacet] = []


FocusedRunRead.model_rebuild()