trustos/backend/app/schemas/schemas.py

from pydantic import BaseModel, EmailStr, field_validator
from typing import Optional, List
from datetime import datetime
from app.models.models import UserRole, FindingSeverity, FindingStatus, FindingCategory, AssetType


# ─── Auth ────────────────────────────────────────────────────────────────────

class LoginRequest(BaseModel):
    email: str
    password: str


class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"
    role: str
    tenant_id: str
    full_name: str


class UserMe(BaseModel):
    id: str
    email: str
    full_name: str
    role: UserRole
    tenant_id: str


# ─── Tenant ───────────────────────────────────────────────────────────────────

class TenantBase(BaseModel):
    name: str
    slug: str
    industry: Optional[str] = None
    size_range: Optional[str] = None
    contact_email: Optional[str] = None


class TenantCreate(TenantBase):
    pass


class TenantOut(TenantBase):
    id: str
    is_active: bool
    created_at: datetime
    model_config = {"from_attributes": True}


# ─── Finding ──────────────────────────────────────────────────────────────────

class FindingOut(BaseModel):
    id: str
    tenant_id: str
    asset_id: Optional[str] = None
    title: str
    severity: FindingSeverity
    status: FindingStatus
    category: FindingCategory
    technical_description: Optional[str] = None
    cve_id: Optional[str] = None
    cvss_score: Optional[float] = None
    affected_component: Optional[str] = None
    ai_summary: Optional[str] = None
    ai_business_impact: Optional[str] = None
    ai_impact_level: Optional[str] = None
    ai_remediation_steps: Optional[str] = None
    ai_fix_priority: Optional[str] = None
    assignee_email: Optional[str] = None
    due_date: Optional[datetime] = None
    is_top_risk: bool
    source: Optional[str] = None
    created_at: datetime
    updated_at: datetime
    model_config = {"from_attributes": True}


class FindingCreate(BaseModel):
    title: str
    severity: FindingSeverity
    category: FindingCategory
    technical_description: Optional[str] = None
    cve_id: Optional[str] = None
    cvss_score: Optional[float] = None
    affected_component: Optional[str] = None
    asset_id: Optional[str] = None
    executive_id: Optional[str] = None
    source: str = "manual"


class FindingStatusUpdate(BaseModel):
    status: FindingStatus
    resolution_note: Optional[str] = None
    assignee_email: Optional[str] = None
    due_date: Optional[datetime] = None


# ─── Risk Score ───────────────────────────────────────────────────────────────

class RiskScoreOut(BaseModel):
    id: str
    tenant_id: str
    score_date: datetime
    overall_score: float
    score_identity: Optional[float] = None
    score_cloud: Optional[float] = None
    score_network: Optional[float] = None
    score_web: Optional[float] = None
    score_credential: Optional[float] = None
    score_digital_footprint: Optional[float] = None
    score_third_party: Optional[float] = None
    critical_count: int
    high_count: int
    medium_count: int
    low_count: int
    model_config = {"from_attributes": True}


# ─── Dashboard ────────────────────────────────────────────────────────────────

class RiskCardData(BaseModel):
    id: str
    title: str
    ai_summary: Optional[str]
    ai_business_impact: Optional[str]
    ai_impact_level: Optional[str]
    ai_fix_priority: Optional[str]
    severity: str
    category: str


class DashboardResponse(BaseModel):
    tenant_name: str
    current_score: float
    previous_score: Optional[float]
    score_delta: Optional[float]
    score_trend: List[dict]  # [{date, score}]
    top_risks: List[RiskCardData]
    open_critical: int
    open_high: int
    open_medium: int
    total_open: int
    baseline_score: Optional[float] = None
    baseline_date: Optional[datetime] = None


# ─── Audit Report ─────────────────────────────────────────────────────────────

class AuditReportOut(BaseModel):
    id: str
    tenant_id: str
    title: str
    report_date: datetime
    baseline_score: Optional[float]
    executive_summary: Optional[str]
    scope_description: Optional[str]
    pdf_path: Optional[str]
    is_baseline: bool
    generated_by: Optional[str]
    created_at: datetime
    model_config = {"from_attributes": True}


class AuditReportCreate(BaseModel):
    title: str
    executive_summary: Optional[str] = None
    scope_description: Optional[str] = None


# ─── Attack Path ──────────────────────────────────────────────────────────────

class AttackPathOut(BaseModel):
    id: str
    finding_id: str
    title: str
    ai_narrative: Optional[str]
    nodes_json: Optional[str]
    edges_json: Optional[str]
    created_at: datetime
    model_config = {"from_attributes": True}


# ─── Authorized Asset ─────────────────────────────────────────────────────────

class AuthorizedAssetCreate(BaseModel):
    value: str
    asset_type: AssetType
    description: Optional[str] = None
    authorized_by: str


class AuthorizedAssetOut(BaseModel):
    id: str
    tenant_id: str
    value: str
    asset_type: AssetType
    description: Optional[str]
    authorized_by: Optional[str]
    authorized_at: datetime
    is_active: bool
    model_config = {"from_attributes": True}