drjones/trustos
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
trustos/backend/app/models/models.py
drjones 463dff883b feat: initial TrustOS platform scaffold 
- FastAPI backend: auth, findings, dashboard, attack paths, footprint, AI translator, risk calculator, PDF report generator
- Next.js frontend: Vault dashboard, login, findings table, finding detail with AI coach, digital footprint, reports
- PostgreSQL data model: tenants, users, assets, findings, risk scores, audit reports, attack paths
- Docker Compose + Dockerfiles for all services
- Demo seed data: Acme Corp with 6 findings and 90-day risk score history
- AI Risk Translator (OpenAI/Anthropic) with plain-English business impact
- Role-based access: executive / it_admin / trustos_admin
- Scope-lock engine: authorization required before any assessment

Stage 1-8 complete: Phase 1 Vault Audit product ready
2026-07-05 09:46:08 +00:00

276 lines
15 KiB
Python
Raw Permalink Blame History

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")
Reference in New Issue View Git Blame Copy Permalink