validation/app/websocket/protocol.py

from pydantic import BaseModel, Field, ConfigDict
from typing import Literal, Optional, List, Dict, Any
from datetime import datetime
from enum import Enum


class MessageType(str, Enum):
    # Client -> Server
    AUDIO_CHUNK = "audio_chunk"
    CONFIG = "config"
    PING = "ping"
    RESET = "reset"

    # Server -> Client
    DETECTION_RESULT = "detection_result"
    THREAT_ALERT = "threat_alert"
    KILL_SWITCH = "kill_switch"
    RISK_ALERT = "risk_alert"
    ALERT = "alert"
    VERIFY_REQUIRED = "verify_required"
    STATS = "stats"
    ERROR = "error"
    PONG = "pong"
    READY = "ready"


class AudioFormat(str, Enum):
    INT16 = "int16"
    FLOAT32 = "float32"


class ClientAudioChunk(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.AUDIO_CHUNK] = MessageType.AUDIO_CHUNK
    data: List[int] = Field(..., description="Audio samples as int16 array")
    sample_rate: int = Field(default=16000, ge=8000, le=48000)
    timestamp_ms: Optional[int] = None
    chunk_index: Optional[int] = None
    format: AudioFormat = AudioFormat.INT16


class ClientConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.CONFIG] = MessageType.CONFIG
    threat_threshold: Optional[float] = Field(default=None, ge=0, le=100)
    consecutive_windows: Optional[int] = Field(default=None, ge=1, le=10)
    enable_vad: Optional[bool] = True
    enable_kill_switch: Optional[bool] = True


class ClientPing(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.PING] = MessageType.PING
    timestamp_ms: int


class ClientReset(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.RESET] = MessageType.RESET


# Union of all client messages
ClientMessage = ClientAudioChunk | ClientConfig | ClientPing | ClientReset


class VadResult(BaseModel):
    is_speech: bool
    energy: float
    snr: float
    confidence: float


class DetectionResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.DETECTION_RESULT] = MessageType.DETECTION_RESULT
    chunk_index: int
    timestamp_ms: int
    probabilities: Dict[str, float]
    predicted_class: int
    predicted_label: str
    confidence: float
    is_bonafide: bool
    threat_score: float
    authenticity_score: float
    inference_ms: float
    vad: VadResult
    audio_hash: str
    # Multi-detector outputs (new)
    replay_score: float = 0.0
    replay_high_confidence: bool = False
    speaker_match: Optional[float] = None
    speaker_similarity: Optional[float] = None
    matched_speaker: Optional[str] = None
    fused_synthetic_probability: float = 0.0
    calibrated_synthetic_probability: float = 0.0
    calibrated_speaker_match: Optional[float] = None
    detector_consistency: float = 0.0
    risk_level: str = "LOW"
    risk_score: float = 0.0
    risk_action: str = "MONITOR"
    decision: str = "ALLOW"
    behavioral_risk: Optional[float] = None
    context_risk: Optional[float] = None
    risk_escalated: bool = False
    detections: Dict[str, Any] = Field(default_factory=dict)


class ThreatAlert(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.THREAT_ALERT] = MessageType.THREAT_ALERT
    threat_score: float
    authenticity_score: float
    predicted_label: str
    confidence: float
    timestamp_ms: int
    chunk_index: int
    audio_hash: str


class KillSwitchMessage(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.KILL_SWITCH] = MessageType.KILL_SWITCH
    action: Literal["TERMINATE_CALL"] = "TERMINATE_CALL"
    reason: str
    threat_score: float
    consecutive_detections: int
    timestamp_ms: int
    audio_hash: str


class StatsMessage(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.STATS] = MessageType.STATS
    total_chunks: int
    avg_inference_ms: float
    kill_switch_triggers: int
    errors: int
    buffer_fill_ratio: float
    uptime_ms: int


class ErrorMessage(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.ERROR] = MessageType.ERROR
    code: str
    message: str
    details: Optional[Dict[str, Any]] = None


class PongMessage(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.PONG] = MessageType.PONG
    timestamp_ms: int
    server_time_ms: int


class RiskAlert(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.RISK_ALERT] = MessageType.RISK_ALERT
    risk_level: str
    risk_score: float
    action: str
    reasons: List[str] = Field(default_factory=list)
    calibrated_synthetic_probability: float
    speaker_match: Optional[float] = None
    replay_probability: float
    consecutive_level: int = 1
    escalated: bool = False
    decision: str = "ALLOW"
    timestamp_ms: int
    chunk_index: int
    audio_hash: str


class VerifyRequired(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.VERIFY_REQUIRED] = MessageType.VERIFY_REQUIRED
    flow_id: str
    call_id: str
    state: str
    risk_level: str
    reason: str = ""
    attempts: int = 0
    prototype: bool = True
    timestamp_ms: int
    message: str = "Verification requested (SIMULATED OTP workflow for the demo)."


class InAppAlert(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal[MessageType.ALERT] = MessageType.ALERT
    id: str
    alert_type: str
    severity: str
    message: str
    call_id: str
    simulated: bool = True
    provider: str = "simulated"
    timestamp_ms: int


class ReadyMessage(BaseModel):
    model_config = ConfigDict(extra="forbid", protected_namespaces=())

    type: Literal[MessageType.READY] = MessageType.READY
    model_info: Dict[str, Any]
    config: Dict[str, Any]


# Union of all server messages
ServerMessage = (
    DetectionResult
    | ThreatAlert
    | KillSwitchMessage
    | RiskAlert
    | VerifyRequired
    | InAppAlert
    | StatsMessage
    | ErrorMessage
    | PongMessage
    | ReadyMessage
)


class AuditLogEntry(BaseModel):
    """Forensic audit log entry (immutable)."""
    timestamp_ms: int
    connection_id: str
    chunk_index: int
    chunk_hash: str
    detected_vocoder: str
    threat_score: float
    authenticity_score: float
    inference_ms: float
    kill_switch_triggered: bool
    kill_switch_reason: str = ""


# Type discriminator for message parsing
MESSAGE_TYPE_FIELD = "type"