import uuid from datetime import datetime from sqlalchemy import String, DateTime, Text, Float, Integer, Boolean, ForeignKey, Enum as SAEnum from sqlalchemy.orm import Mapped, mapped_column, relationship from sqlalchemy.dialects.postgresql import UUID as PGUUID from app.db.session import Base import enum class UserRole(str, enum.Enum): executive = "executive" it_admin = "it_admin" trustos_admin = "trustos_admin" class FindingSeverity(str, enum.Enum): critical = "critical" high = "high" medium = "medium" low = "low" info = "info" class FindingStatus(str, enum.Enum): open = "open" in_progress = "in_progress" resolved = "resolved" verified = "verified" accepted_risk = "accepted_risk" class FindingCategory(str, enum.Enum): external_exposure = "external_exposure" cloud_posture = "cloud_posture" credential_exposure = "credential_exposure" digital_footprint = "digital_footprint" web_application = "web_application" network = "network" identity = "identity" third_party = "third_party" compliance = "compliance" other = "other" class AssetType(str, enum.Enum): domain = "domain" ip_address = "ip_address" cloud_resource = "cloud_resource" email_account = "email_account" executive = "executive" web_application = "web_application" api_endpoint = "api_endpoint" repository = "repository" # ───────────────────────────────────────── # Tenant (client organization) # ───────────────────────────────────────── class Tenant(Base): __tablename__ = "tenants" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) name: Mapped[str] = mapped_column(String(255), nullable=False) slug: Mapped[str] = mapped_column(String(100), unique=True, nullable=False) industry: Mapped[str | None] = mapped_column(String(100)) size_range: Mapped[str | None] = mapped_column(String(50)) # e.g. "50-200" contact_email: Mapped[str | None] = mapped_column(String(255)) is_active: Mapped[bool] = mapped_column(Boolean, default=True) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) users: Mapped[list["User"]] = relationship("User", back_populates="tenant", lazy="select") assets: Mapped[list["Asset"]] = relationship("Asset", back_populates="tenant", lazy="select") findings: Mapped[list["Finding"]] = relationship("Finding", back_populates="tenant", lazy="select") risk_scores: Mapped[list["RiskScore"]] = relationship("RiskScore", back_populates="tenant", lazy="select") audit_reports: Mapped[list["AuditReport"]] = relationship("AuditReport", back_populates="tenant", lazy="select") executives: Mapped[list["Executive"]] = relationship("Executive", back_populates="tenant", lazy="select") authorized_assets: Mapped[list["AuthorizedAsset"]] = relationship("AuthorizedAsset", back_populates="tenant", lazy="select") # ───────────────────────────────────────── # User # ───────────────────────────────────────── class User(Base): __tablename__ = "users" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False) hashed_password: Mapped[str] = mapped_column(String(255), nullable=False) full_name: Mapped[str] = mapped_column(String(255), nullable=False) role: Mapped[UserRole] = mapped_column(SAEnum(UserRole), default=UserRole.it_admin) is_active: Mapped[bool] = mapped_column(Boolean, default=True) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) last_login: Mapped[datetime | None] = mapped_column(DateTime) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="users") # ───────────────────────────────────────── # Authorized Asset Scope (scope-lock) # ───────────────────────────────────────── class AuthorizedAsset(Base): __tablename__ = "authorized_assets" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) value: Mapped[str] = mapped_column(String(500), nullable=False) # domain, CIDR, cloud account ID asset_type: Mapped[AssetType] = mapped_column(SAEnum(AssetType), nullable=False) description: Mapped[str | None] = mapped_column(Text) scope_agreement_ref: Mapped[str | None] = mapped_column(String(255)) # file path or reference authorized_by: Mapped[str | None] = mapped_column(String(255)) # email of authorizing contact authorized_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) is_active: Mapped[bool] = mapped_column(Boolean, default=True) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="authorized_assets") # ───────────────────────────────────────── # Asset (discovered asset record) # ───────────────────────────────────────── class Asset(Base): __tablename__ = "assets" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) name: Mapped[str] = mapped_column(String(500), nullable=False) asset_type: Mapped[AssetType] = mapped_column(SAEnum(AssetType), nullable=False) value: Mapped[str] = mapped_column(String(500), nullable=False) # FQDN, IP, ARN, etc. owner_name: Mapped[str | None] = mapped_column(String(255)) owner_email: Mapped[str | None] = mapped_column(String(255)) tags: Mapped[str | None] = mapped_column(Text) # JSON array of tags notes: Mapped[str | None] = mapped_column(Text) is_active: Mapped[bool] = mapped_column(Boolean, default=True) first_seen: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) last_seen: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="assets") findings: Mapped[list["Finding"]] = relationship("Finding", back_populates="asset", lazy="select") # ───────────────────────────────────────── # Executive (for digital footprint tracking) # ───────────────────────────────────────── class Executive(Base): __tablename__ = "executives" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) full_name: Mapped[str] = mapped_column(String(255), nullable=False) title: Mapped[str | None] = mapped_column(String(255)) corporate_email: Mapped[str | None] = mapped_column(String(255)) is_enrolled: Mapped[bool] = mapped_column(Boolean, default=True) # explicit consent notes: Mapped[str | None] = mapped_column(Text) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="executives") # ───────────────────────────────────────── # Finding # ───────────────────────────────────────── class Finding(Base): __tablename__ = "findings" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) asset_id: Mapped[str | None] = mapped_column(String(36), ForeignKey("assets.id")) executive_id: Mapped[str | None] = mapped_column(String(36), ForeignKey("executives.id")) title: Mapped[str] = mapped_column(String(500), nullable=False) severity: Mapped[FindingSeverity] = mapped_column(SAEnum(FindingSeverity), nullable=False) status: Mapped[FindingStatus] = mapped_column(SAEnum(FindingStatus), default=FindingStatus.open) category: Mapped[FindingCategory] = mapped_column(SAEnum(FindingCategory), nullable=False) # Technical fields (IT Admin view) technical_description: Mapped[str | None] = mapped_column(Text) cve_id: Mapped[str | None] = mapped_column(String(30)) cvss_score: Mapped[float | None] = mapped_column(Float) affected_component: Mapped[str | None] = mapped_column(String(500)) evidence: Mapped[str | None] = mapped_column(Text) # raw evidence # AI-translated fields (Executive view) ai_summary: Mapped[str | None] = mapped_column(Text) # plain-English summary ai_business_impact: Mapped[str | None] = mapped_column(Text) ai_impact_level: Mapped[str | None] = mapped_column(String(20)) # Low/Medium/High/Critical ai_remediation_steps: Mapped[str | None] = mapped_column(Text) ai_fix_priority: Mapped[str | None] = mapped_column(String(20)) # urgent/soon/planned ai_generated_at: Mapped[datetime | None] = mapped_column(DateTime) # Remediation tracking assignee_email: Mapped[str | None] = mapped_column(String(255)) due_date: Mapped[datetime | None] = mapped_column(DateTime) resolution_note: Mapped[str | None] = mapped_column(Text) resolved_at: Mapped[datetime | None] = mapped_column(DateTime) verified_at: Mapped[datetime | None] = mapped_column(DateTime) # Metadata source: Mapped[str | None] = mapped_column(String(100)) # "manual", "scanner", "hibp", "cve_monitor" is_top_risk: Mapped[bool] = mapped_column(Boolean, default=False) # flagged as a Top 3 risk created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="findings") asset: Mapped["Asset"] = relationship("Asset", back_populates="findings") attack_paths: Mapped[list["AttackPath"]] = relationship("AttackPath", back_populates="finding", lazy="select") # ───────────────────────────────────────── # Risk Score (daily snapshot) # ───────────────────────────────────────── class RiskScore(Base): __tablename__ = "risk_scores" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) score_date: Mapped[datetime] = mapped_column(DateTime, nullable=False) overall_score: Mapped[float] = mapped_column(Float, nullable=False) # Category scores score_identity: Mapped[float | None] = mapped_column(Float) score_cloud: Mapped[float | None] = mapped_column(Float) score_network: Mapped[float | None] = mapped_column(Float) score_web: Mapped[float | None] = mapped_column(Float) score_credential: Mapped[float | None] = mapped_column(Float) score_digital_footprint: Mapped[float | None] = mapped_column(Float) score_third_party: Mapped[float | None] = mapped_column(Float) # Counts at time of snapshot critical_count: Mapped[int] = mapped_column(Integer, default=0) high_count: Mapped[int] = mapped_column(Integer, default=0) medium_count: Mapped[int] = mapped_column(Integer, default=0) low_count: Mapped[int] = mapped_column(Integer, default=0) notes: Mapped[str | None] = mapped_column(Text) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="risk_scores") # ───────────────────────────────────────── # Attack Path # ───────────────────────────────────────── class AttackPath(Base): __tablename__ = "attack_paths" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) finding_id: Mapped[str] = mapped_column(String(36), ForeignKey("findings.id"), nullable=False) title: Mapped[str] = mapped_column(String(500), nullable=False) ai_narrative: Mapped[str | None] = mapped_column(Text) nodes_json: Mapped[str | None] = mapped_column(Text) # JSON: [{id, label, type, risk_level}] edges_json: Mapped[str | None] = mapped_column(Text) # JSON: [{source, target}] created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) finding: Mapped["Finding"] = relationship("Finding", back_populates="attack_paths") # ───────────────────────────────────────── # Audit Report (Phase 1 deliverable) # ───────────────────────────────────────── class AuditReport(Base): __tablename__ = "audit_reports" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4())) tenant_id: Mapped[str] = mapped_column(String(36), ForeignKey("tenants.id"), nullable=False) title: Mapped[str] = mapped_column(String(500), nullable=False) report_date: Mapped[datetime] = mapped_column(DateTime, nullable=False, default=datetime.utcnow) baseline_score: Mapped[float | None] = mapped_column(Float) executive_summary: Mapped[str | None] = mapped_column(Text) scope_description: Mapped[str | None] = mapped_column(Text) key_findings_json: Mapped[str | None] = mapped_column(Text) # JSON snapshot of top findings pdf_path: Mapped[str | None] = mapped_column(String(500)) is_baseline: Mapped[bool] = mapped_column(Boolean, default=True) generated_by: Mapped[str | None] = mapped_column(String(255)) # admin email created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) tenant: Mapped["Tenant"] = relationship("Tenant", back_populates="audit_reports")