rengine/shared/models/target.py

import uuid
from datetime import datetime
from typing import TYPE_CHECKING

from pydantic import BaseModel, field_validator
from sqlmodel import Field, Relationship, SQLModel, UniqueConstraint

from shared.definitions.rescan import MAX_RUN_ASSETS
from shared.enums.target import TargetType
from shared.enums.task_status import TaskStatus
from shared.models.bgp_summary import BgpSummaryRead, TargetBgpSummary
from shared.models.dns import DnsLookup, DnsLookupSummary
from shared.models.organization import Organization, OrganizationSummary
from shared.models.tag import TagSummary, TargetTag
from shared.models.whois import WhoisRecord, WhoisRecordSummary
from shared.utils.datetime import utc_now
from shared.utils.validation import clean_name

if TYPE_CHECKING:
    from shared.models.tag import Tag


MAX_TAG_LEN = 50
MAX_ORG_LEN = 100


def _clean_labels(values: list[str] | None, max_len: int) -> list[str] | None:
    """Drop blank entries and cap each label's length."""
    if values is None:
        return None
    return [clean_name(v, max_len=max_len) for v in values if (v or "").strip()]


class TargetValidationRequest(BaseModel):
    target_value: str


class TargetValidationResponse(BaseModel):
    valid: bool
    target_type: TargetType | None
    error: str | None
    target_value: str


class TargetOrganization(SQLModel, table=True):
    __tablename__ = "target_organizations"

    target_id: uuid.UUID = Field(
        foreign_key="targets.id", primary_key=True, ondelete="CASCADE"
    )
    organization_id: uuid.UUID = Field(
        foreign_key="organizations.id", primary_key=True, ondelete="CASCADE"
    )


class TargetBase(SQLModel):
    target_value: str = Field(max_length=500)
    display_name: str | None = Field(default=None, max_length=200)


class Target(TargetBase, table=True):
    __tablename__ = "targets"
    __table_args__ = (
        UniqueConstraint("target_value", "project_id", name="uq_target_value_project"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    target_type: TargetType
    project_id: uuid.UUID = Field(foreign_key="projects.id", index=True)
    created_at: datetime = Field(default_factory=utc_now)
    updated_at: datetime = Field(default_factory=utc_now)
    created_by: uuid.UUID = Field(foreign_key="users.id")

    whois_status: TaskStatus = Field(default=TaskStatus.PENDING, index=True)
    whois_error: str | None = Field(default=None, max_length=1000)
    whois_record: WhoisRecord | None = Relationship(
        sa_relationship_kwargs={"lazy": "selectin"},
    )
    whois_record_id: uuid.UUID | None = Field(
        default=None, foreign_key="whois_records.id", index=True, ondelete="SET NULL"
    )

    bgp_status: TaskStatus = Field(default=TaskStatus.PENDING, index=True)
    bgp_summary: TargetBgpSummary | None = Relationship(
        sa_relationship_kwargs={"lazy": "selectin", "passive_deletes": True},
    )

    seed_scans: bool = Field(default=True)

    dns_status: TaskStatus = Field(default=TaskStatus.PENDING, index=True)
    dns_error: str | None = Field(default=None, max_length=1000)
    dns_lookup: DnsLookup | None = Relationship(
        sa_relationship_kwargs={
            "lazy": "selectin",
            "foreign_keys": "Target.dns_lookup_id",
        },
    )
    dns_lookup_id: uuid.UUID | None = Field(
        default=None, foreign_key="dns_lookups.id", index=True, ondelete="SET NULL"
    )

    organizations: list["Organization"] = Relationship(
        link_model=TargetOrganization,
        sa_relationship_kwargs={"lazy": "selectin"},
    )

    tags: list["Tag"] = Relationship(
        back_populates="targets",
        link_model=TargetTag,
        sa_relationship_kwargs={"lazy": "selectin"},
    )


class TargetCreate(TargetBase):
    project_slug: str
    organization_names: list[str] = Field(default_factory=list)
    tag_names: list[str] = Field(default_factory=list)
    seeds: list[str] = Field(default_factory=list, max_length=MAX_RUN_ASSETS)

    @field_validator("tag_names", check_fields=False)
    @classmethod
    def _clean_tags(cls, v):
        return _clean_labels(v, MAX_TAG_LEN)

    @field_validator("organization_names", check_fields=False)
    @classmethod
    def _clean_orgs(cls, v):
        return _clean_labels(v, MAX_ORG_LEN)


# bulk targets create/request/response models
class TargetBulkCreate(BaseModel):
    project_slug: str
    targets: list[str] = Field(
        ...,
        min_length=1,
        max_length=1000,
        description="Target values to import. At most 1000.",
    )
    organization_names: list[str] = Field(default_factory=list)
    tag_names: list[str] = Field(default_factory=list)
    seeds: list[str] = Field(default_factory=list, max_length=MAX_RUN_ASSETS)

    @field_validator("tag_names", "tags", check_fields=False)
    @classmethod
    def _clean_tags(cls, v):
        return _clean_labels(v, MAX_TAG_LEN)

    @field_validator("organization_names", "organizations", check_fields=False)
    @classmethod
    def _clean_orgs(cls, v):
        return _clean_labels(v, MAX_ORG_LEN)


class TargetImportResult(BaseModel):
    # single target import result
    target_value: str
    success: bool
    target_type: TargetType | None = None
    target_id: uuid.UUID | None = None
    error: str | None = None


class TargetBulkCreateResponse(BaseModel):
    total: int
    imported: int
    failed: int
    skipped_duplicates: int
    results: list[TargetImportResult]


class TargetUpdate(SQLModel):
    display_name: str | None = Field(default=None, max_length=200)
    organization_names: list[str] | None = None
    tag_names: list[str] | None = None
    seed_scans: bool | None = None

    @field_validator("tag_names", "tags", check_fields=False)
    @classmethod
    def _clean_tags(cls, v):
        return _clean_labels(v, MAX_TAG_LEN)

    @field_validator("organization_names", "organizations", check_fields=False)
    @classmethod
    def _clean_orgs(cls, v):
        return _clean_labels(v, MAX_ORG_LEN)


class TargetRead(TargetBase):
    id: uuid.UUID
    target_type: TargetType
    project_id: uuid.UUID
    created_at: datetime
    updated_at: datetime
    created_by: uuid.UUID
    whois_status: TaskStatus
    whois_error: str | None
    whois_record_id: uuid.UUID | None
    whois: WhoisRecordSummary | None = None
    bgp_status: TaskStatus
    bgp: BgpSummaryRead | None = None
    dns_status: TaskStatus = TaskStatus.PENDING
    dns_error: str | None = None
    dns_lookup_id: uuid.UUID | None = None
    dns: DnsLookupSummary | None = None
    organizations: list[OrganizationSummary] = Field(default_factory=list)
    tags: list[TagSummary] = Field(default_factory=list)
    seed_scans: bool = True
    seed_count: int | None = None


# imports for both json and csv support
class TargetImportItem(BaseModel):
    target_value: str
    tags: list[str] = Field(default_factory=list)
    organizations: list[str] = Field(default_factory=list)
    display_name: str | None = None
    seeds: list[str] = Field(default_factory=list, max_length=MAX_RUN_ASSETS)

    @field_validator("tag_names", "tags", check_fields=False)
    @classmethod
    def _clean_tags(cls, v):
        return _clean_labels(v, MAX_TAG_LEN)

    @field_validator("organization_names", "organizations", check_fields=False)
    @classmethod
    def _clean_orgs(cls, v):
        return _clean_labels(v, MAX_ORG_LEN)


class TargetImportRequest(BaseModel):
    project_slug: str
    targets: list[TargetImportItem] = Field(
        ...,
        min_length=1,
        max_length=500,
        description="Targets to import. At most 500.",
    )
    organization_names: list[str] = Field(
        default_factory=list,
        description="Organizations added to every imported target.",
    )
    tag_names: list[str] = Field(
        default_factory=list,
        description="Tags added to every imported target.",
    )

    @field_validator("tag_names", "tags", check_fields=False)
    @classmethod
    def _clean_tags(cls, v):
        return _clean_labels(v, MAX_TAG_LEN)

    @field_validator("organization_names", "organizations", check_fields=False)
    @classmethod
    def _clean_orgs(cls, v):
        return _clean_labels(v, MAX_ORG_LEN)