path: root/src/nvd/models.py

"""Pydantic models for NVD API responses."""

from datetime import datetime
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field


# Common models
class LangString(BaseModel):
    """Multilingual string."""

    lang: str
    value: str


# CVSS Models
class CVSSMetricV2(BaseModel):
    """CVSS v2.0 metrics."""

    source: str
    type: str
    cvssData: Dict[str, Any]
    baseSeverity: Optional[str] = None
    exploitabilityScore: Optional[float] = None
    impactScore: Optional[float] = None
    acInsufInfo: Optional[bool] = None
    obtainAllPrivilege: Optional[bool] = None
    obtainUserPrivilege: Optional[bool] = None
    obtainOtherPrivilege: Optional[bool] = None
    userInteractionRequired: Optional[bool] = None


class CVSSMetricV3(BaseModel):
    """CVSS v3.x metrics."""

    source: str
    type: str
    cvssData: Dict[str, Any]
    exploitabilityScore: Optional[float] = None
    impactScore: Optional[float] = None


class CVSSMetricV4(BaseModel):
    """CVSS v4.0 metrics."""

    source: str
    type: str
    cvssData: Dict[str, Any]


class Metrics(BaseModel):
    """CVSS metrics container."""

    cvssMetricV2: Optional[List[CVSSMetricV2]] = None
    cvssMetricV30: Optional[List[CVSSMetricV3]] = None
    cvssMetricV31: Optional[List[CVSSMetricV3]] = None
    cvssMetricV40: Optional[List[CVSSMetricV4]] = None


# CPE Models
class CPEMatch(BaseModel):
    """CPE match criteria."""

    vulnerable: bool
    criteria: str
    matchCriteriaId: str
    versionStartIncluding: Optional[str] = None
    versionStartExcluding: Optional[str] = None
    versionEndIncluding: Optional[str] = None
    versionEndExcluding: Optional[str] = None


class Node(BaseModel):
    """Configuration node."""

    operator: str
    negate: bool
    cpeMatch: List[CPEMatch]


class Configuration(BaseModel):
    """CVE configuration."""

    nodes: List[Node]


# Weakness Models
class WeaknessDescription(BaseModel):
    """CWE description."""

    lang: str
    value: str


class Weakness(BaseModel):
    """CWE weakness."""

    source: str
    type: str
    description: List[WeaknessDescription]


# Reference Models
class Reference(BaseModel):
    """External reference."""

    url: str
    source: Optional[str] = None
    tags: Optional[List[str]] = None


# CVE Models
class CVEData(BaseModel):
    """Core CVE data."""

    id: str = Field(alias="cveId")
    sourceIdentifier: str
    published: datetime
    lastModified: datetime
    vulnStatus: str
    cveTags: Optional[List[Dict[str, str]]] = None
    descriptions: List[LangString]
    metrics: Optional[Metrics] = None
    weaknesses: Optional[List[Weakness]] = None
    configurations: Optional[List[Configuration]] = None
    references: List[Reference]
    vendorComments: Optional[List[Dict[str, Any]]] = None

    class Config:
        populate_by_name = True

    @property
    def description(self) -> str:
        """Get English description."""
        for desc in self.descriptions:
            if desc.lang == "en":
                return desc.value
        return self.descriptions[0].value if self.descriptions else ""

    @property
    def cvss_v3_score(self) -> Optional[float]:
        """Get primary CVSS v3.x base score."""
        if self.metrics:
            if self.metrics.cvssMetricV31:
                return self.metrics.cvssMetricV31[0].cvssData.get("baseScore")
            if self.metrics.cvssMetricV30:
                return self.metrics.cvssMetricV30[0].cvssData.get("baseScore")
        return None

    @property
    def cvss_v2_score(self) -> Optional[float]:
        """Get CVSS v2.0 base score."""
        if self.metrics and self.metrics.cvssMetricV2:
            return self.metrics.cvssMetricV2[0].cvssData.get("baseScore")
        return None


class CVEItem(BaseModel):
    """CVE item wrapper."""

    cve: CVEData


class CVEResponse(BaseModel):
    """CVE API response."""

    resultsPerPage: int
    startIndex: int
    totalResults: int
    format: str
    version: str
    timestamp: datetime
    vulnerabilities: List[CVEItem]


# CVE Change History Models
class ChangeDetail(BaseModel):
    """Individual change detail."""

    action: Optional[str] = None
    type: Optional[str] = None
    oldValue: Optional[str] = None
    newValue: Optional[str] = None


class CVEChange(BaseModel):
    """CVE change event."""

    cveId: str
    eventName: str
    cveChangeId: str
    sourceIdentifier: str
    created: datetime
    details: Optional[List[ChangeDetail]] = None


class CVEChangeResponse(BaseModel):
    """CVE Change History API response."""

    resultsPerPage: int
    startIndex: int
    totalResults: int
    format: str
    version: str
    timestamp: datetime
    cveChanges: List[CVEChange]


# CPE Models
class CPETitle(BaseModel):
    """CPE title in different languages."""

    title: str
    lang: str


class CPERef(BaseModel):
    """CPE reference."""

    ref: str
    type: Optional[str] = None


class CPEData(BaseModel):
    """CPE product data."""

    cpeName: str
    cpeNameId: str
    created: datetime
    lastModified: datetime
    deprecated: bool = False
    titles: Optional[List[CPETitle]] = None
    refs: Optional[List[CPERef]] = None
    deprecatedBy: Optional[List[Dict[str, str]]] = None

    @property
    def title(self) -> str:
        """Get English title."""
        if self.titles:
            for t in self.titles:
                if t.lang == "en":
                    return t.title
            return self.titles[0].title
        return ""


class CPEItem(BaseModel):
    """CPE item wrapper from API response."""

    cpe: CPEData


class CPEResponse(BaseModel):
    """CPE API response."""

    resultsPerPage: int
    startIndex: int
    totalResults: int
    format: str
    version: str
    timestamp: datetime
    products: List[CPEItem]


# CPE Match Criteria Models
class CPEMatchString(BaseModel):
    """CPE match string data."""

    matchCriteriaId: str
    criteria: str
    lastModified: datetime
    cpeLastModified: Optional[datetime] = None
    created: datetime
    status: str
    matches: Optional[List[Dict[str, str]]] = None
    versionStartIncluding: Optional[str] = None
    versionStartExcluding: Optional[str] = None
    versionEndIncluding: Optional[str] = None
    versionEndExcluding: Optional[str] = None


class CPEMatchResponse(BaseModel):
    """CPE Match Criteria API response."""

    resultsPerPage: int
    startIndex: int
    totalResults: int
    format: str
    version: str
    timestamp: datetime
    matchStrings: List[CPEMatchString]


# Source Models
class SourceData(BaseModel):
    """Data source organization."""

    name: str
    contactEmail: str
    sourceIdentifiers: List[str]
    lastModified: datetime
    created: datetime


class SourceResponse(BaseModel):
    """Source API response."""

    resultsPerPage: int
    startIndex: int
    totalResults: int
    format: str
    version: str
    timestamp: datetime
    sources: List[SourceData]