mihir2007/Cyber-Risk

1"""2models.py3---------4SQLAlchemy 2.0 declarative ORM models for the CRQ platform.5 6Tables7------8Asset            : enterprise asset inventory (hosts, gateways, databases...)9Vulnerability    : CVEs discovered on a given asset (many-to-one -> Asset)10SecurityControl  : catalogue of candidate security investments11SimulationRun    : historical audit log of every Monte Carlo / optimizer run12NetworkEdge      : attack graph topology edges persisting lateral movement paths13DataLineageFlow  : transitive data flows across nodes tracking provenance and leaks14"""15 16from __future__ import annotations17 18import datetime as dt19import json20from typing import Any, Optional21 22from sqlalchemy import (23    Boolean,24    DateTime,25    Float,26    ForeignKey,27    Index,28    Integer,29    String,30    Text,31)32from sqlalchemy.dialects.postgresql import JSONB33from sqlalchemy.orm import Mapped, mapped_column, relationship34from sqlalchemy.types import TypeDecorator35 36from database import Base37 38 39def _utcnow() -> dt.datetime:40    """Returns current UTC timestamp with timezone awareness."""41    return dt.datetime.now(dt.timezone.utc)42 43 44class SafeJSONB(TypeDecorator):45    """46    PostgreSQL native JSONB column type with transparent Python serialization.47 48    Accepts either Python objects (dicts, lists) or JSON strings on write,49    and returns a valid JSON string on read for compatibility with Pydantic50    schemas and REST consumers while storing native JSONB in PostgreSQL.51    """52 53    impl = JSONB54    cache_ok = True55 56    def load_dialect_impl(self, dialect: Any) -> Any:57        if dialect.name == "postgresql":58            return dialect.type_descriptor(JSONB())59        return dialect.type_descriptor(Text())60 61    def process_bind_param(self, value: Any, dialect: Any) -> Any:62        if value is None:63            return None64        if dialect.name == "postgresql":65            if isinstance(value, str):66                try:67                    return json.loads(value)68                except (ValueError, TypeError):69                    return value70            return value71        else:72            if not isinstance(value, str):73                return json.dumps(value)74            return value75 76    def process_result_value(self, value: Any, dialect: Any) -> Any:77        if value is None:78            return None79        if isinstance(value, (dict, list)):80            return json.dumps(value)81        return str(value)82 83 84class Asset(Base):85    """An enterprise IT/OT asset that carries business and cyber risk."""86 87    __tablename__ = "assets"88 89    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)90    hostname: Mapped[str] = mapped_column(String(128), unique=True, index=True, nullable=False)91 92    # e.g. 'Core Banking Switch', 'UPI Gateway', 'Customer DB', 'Internal ERP',93    #      'Staff Portal', 'Office POS'94    asset_type: Mapped[str] = mapped_column(String(64), nullable=False)95 96    # 'Critical' | 'Medium' | 'Low'97    tier: Mapped[str] = mapped_column(String(16), nullable=False, index=True)98 99    business_unit: Mapped[str] = mapped_column(String(64), nullable=False)100 101    # Revenue lost per minute of downtime, in INR (₹).102    revenue_per_minute: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)103 104    pii_records_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)105    financial_records_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)106 107    is_rbi_regulated: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)108    is_sebi_regulated: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)109 110    # Graph-topology hint: number of network hops from public internet edge111    network_hops_from_internet: Mapped[int] = mapped_column(Integer, nullable=False, default=3)112 113    # Classification taxonomy: 'Core IT', 'AI/ML Model Pipeline', 'Third-Party SaaS',114    # 'Third-Party Vendor API', 'Cloud Infrastructure', 'Data Repository'115    classification_type: Mapped[str] = mapped_column(String(64), default="Internal IT", nullable=False)116 117    # ISO compliance applicability flags118    is_iso27001_regulated: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)119    is_iso42001_regulated: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)120 121    # Third-party vendor risk fields122    is_third_party: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)123    vendor_name: Mapped[Optional[str]] = mapped_column(String(128), nullable=True)124    vendor_risk_tier: Mapped[Optional[str]] = mapped_column(String(32), nullable=True)  # 'Tier-1 Critical', 'Tier-2 High', 'Tier-3 Standard'125    soc2_attestation: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)126 127    # Data residency and cross-border transfer mechanics128    data_residency_country: Mapped[str] = mapped_column(String(8), default="IN", nullable=False)129    cross_border_transfer_enabled: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)130    destination_countries: Mapped[Optional[str]] = mapped_column(String(256), nullable=True)131    transfer_legal_mechanism: Mapped[Optional[str]] = mapped_column(String(64), default="None", nullable=True)132    is_rbi_localization_compliant: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)133 134    vulnerabilities: Mapped[list[Vulnerability]] = relationship(135        "Vulnerability",136        back_populates="asset",137        cascade="all, delete-orphan",138        lazy="selectin",139    )140 141    def __repr__(self) -> str:142        return f"<Asset id={self.id} hostname={self.hostname!r} tier={self.tier}>"143 144 145class Vulnerability(Base):146    """A single CVE observed on a given asset."""147 148    __tablename__ = "vulnerabilities"149 150    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)151    asset_id: Mapped[int] = mapped_column(152        Integer, ForeignKey("assets.id", ondelete="CASCADE"), nullable=False, index=True153    )154 155    cve_id: Mapped[str] = mapped_column(String(32), index=True, nullable=False)156    cvss_score: Mapped[float] = mapped_column(Float, nullable=False)  # 0.0 - 10.0157    epss_score: Mapped[float] = mapped_column(Float, nullable=False)  # 0.0 - 1.0158    cisa_kev: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)159    patch_available: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)160    is_patched: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)161 162    asset: Mapped[Asset] = relationship("Asset", back_populates="vulnerabilities")163 164    def __repr__(self) -> str:165        return f"<Vulnerability {self.cve_id} asset_id={self.asset_id} cvss={self.cvss_score}>"166 167 168class SecurityControl(Base):169    """A candidate security investment considered by the budget optimizer."""170 171    __tablename__ = "security_controls"172 173    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)174    code: Mapped[str] = mapped_column(String(32), unique=True, index=True, nullable=False)175    name: Mapped[str] = mapped_column(String(128), nullable=False)176 177    # 'All' | 'Critical' | 'Medium'178    target_tier: Mapped[str] = mapped_column(String(16), nullable=False, default="All", index=True)179 180    cost_inr: Mapped[float] = mapped_column(Float, nullable=False)181    likelihood_reduction: Mapped[float] = mapped_column(Float, nullable=False)  # 0.0 - 1.0182    is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)183 184    def __repr__(self) -> str:185        return f"<SecurityControl {self.code} cost={self.cost_inr} reduction={self.likelihood_reduction}>"186 187 188class SimulationRun(Base):189    """Historical audit record of an enterprise risk quantification run."""190 191    __tablename__ = "simulation_runs"192 193    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)194    timestamp: Mapped[dt.datetime] = mapped_column(195        DateTime(timezone=True), default=_utcnow, nullable=False, index=True196    )197 198    total_eal_inr: Mapped[float] = mapped_column(Float, nullable=False)199    var_95_inr: Mapped[float] = mapped_column(Float, nullable=False)200    var_99_inr: Mapped[float] = mapped_column(Float, nullable=False)201 202    allocated_budget_inr: Mapped[Optional[float]] = mapped_column(Float, nullable=True)203    selected_controls_json: Mapped[Optional[str]] = mapped_column(SafeJSONB, nullable=True)204 205    def __repr__(self) -> str:206        return f"<SimulationRun id={self.id} EAL={self.total_eal_inr:.2f} @ {self.timestamp}>"207 208 209class NetworkEdge(Base):210    """211    Topology edge representing a network route or lateral attack vector212    between infrastructure nodes (e.g. perimeter gateway, switch, host).213    """214 215    __tablename__ = "network_edges"216 217    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)218    source_node: Mapped[str] = mapped_column(String(128), nullable=False, index=True)219    target_node: Mapped[str] = mapped_column(String(128), nullable=False, index=True)220    weight: Mapped[float] = mapped_column(Float, nullable=False, default=1.0)221    protocol: Mapped[str] = mapped_column(String(32), nullable=False, default="HTTPS")222    is_segmented: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)223 224    # Consent and delegation flags for data provenance & lineage225    is_authorized_flow: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)226    consent_recorded: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)227    flow_type: Mapped[str] = mapped_column(String(64), default="Direct API", nullable=False)  # 'Direct API', 'Sub-Processor Delegation', 'Background Sync', 'Unauthorized Pivot'228    data_tags: Mapped[Optional[str]] = mapped_column(String(256), nullable=True)  # e.g. "Asset-A:PII,Asset-A:FINANCIAL"229 230    __table_args__ = (231        Index("ix_network_edges_source_target", "source_node", "target_node"),232    )233 234    def __repr__(self) -> str:235        return (236            f"<NetworkEdge id={self.id} {self.source_node} -> {self.target_node} "237            f"weight={self.weight} proto={self.protocol} segmented={self.is_segmented} "238            f"auth={self.is_authorized_flow} consent={self.consent_recorded}>"239        )240 241 242class DataLineageFlow(Base):243    """244    Represents transitive data flows across nodes (A -> B -> C), tracking245    originating asset provenance, consent, and potential unauthorized sub-processing.246    """247 248    __tablename__ = "data_lineage_flows"249 250    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)251    origin_asset_id: Mapped[int] = mapped_column(252        Integer, ForeignKey("assets.id", ondelete="CASCADE"), nullable=False, index=True253    )254    intermediary_asset_id: Mapped[Optional[int]] = mapped_column(255        Integer, ForeignKey("assets.id", ondelete="CASCADE"), nullable=True, index=True256    )257    destination_asset_id: Mapped[int] = mapped_column(258        Integer, ForeignKey("assets.id", ondelete="CASCADE"), nullable=False, index=True259    )260 261    is_authorized: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)262    has_user_consent: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)263    records_exposed_pii: Mapped[int] = mapped_column(Integer, default=0, nullable=False)264    records_exposed_financial: Mapped[int] = mapped_column(Integer, default=0, nullable=False)265    detection_status: Mapped[str] = mapped_column(266        String(32), default="Legitimate Flow", nullable=False267    )  # 'Legitimate Flow', 'Unauthorized Sub-Processing', 'Shadow Exfiltration Vector'268 269    origin_asset: Mapped[Asset] = relationship("Asset", foreign_keys=[origin_asset_id])270    intermediary_asset: Mapped[Optional[Asset]] = relationship("Asset", foreign_keys=[intermediary_asset_id])271    destination_asset: Mapped[Asset] = relationship("Asset", foreign_keys=[destination_asset_id])272 273    def __repr__(self) -> str:274        return (275            f"<DataLineageFlow id={self.id} origin={self.origin_asset_id} -> "276            f"inter={self.intermediary_asset_id} -> dest={self.destination_asset_id} "277            f"status='{self.detection_status}'>"278        )279 280