| 1 | import re |
| 2 | from datetime import datetime |
| 3 | from enum import Enum |
| 4 | from typing import List |
| 5 | from typing import Optional |
| 6 | |
| 7 | from pydantic import BaseModel |
| 8 | from pydantic import Field |
| 9 | from pydantic import field_validator |
| 10 | from pydantic import model_validator |
| 11 | |
| 12 | # --- Enums --- |
| 13 | |
| 14 | |
| 15 | class JobStatus(str, Enum): |
| 16 | PENDING = "pending" |
| 17 | RUNNING = "running" |
| 18 | COMPLETED = "completed" |
| 19 | FAILED = "failed" |
| 20 | |
| 21 | |
| 22 | class TriggeredBy(str, Enum): |
| 23 | SCHEDULED = "scheduled" |
| 24 | MANUAL = "manual" |
| 25 | WEBHOOK = "webhook" |
| 26 | |
| 27 | |
| 28 | class SeverityAssessment(str, Enum): |
| 29 | CRITICAL = "Critical" |
| 30 | HIGH = "High" |
| 31 | MEDIUM = "Medium" |
| 32 | LOW = "Low" |
| 33 | INFORMATIONAL = "Informational" |
| 34 | |
| 35 | |
| 36 | class IocType(str, Enum): |
| 37 | IP = "ip" |
| 38 | DOMAIN = "domain" |
| 39 | HASH = "hash" |
| 40 | PROCESS = "process" |
| 41 | URL = "url" |
| 42 | USER = "user" |
| 43 | COMMAND = "command" |
| 44 | |
| 45 | |
| 46 | class VtVerdict(str, Enum): |
| 47 | MALICIOUS = "malicious" |
| 48 | SUSPICIOUS = "suspicious" |
| 49 | CLEAN = "clean" |
| 50 | UNKNOWN = "unknown" |
| 51 | |
| 52 | |
| 53 | class OverallVerdict(str, Enum): |
| 54 | UP = "up" |
| 55 | DOWN = "down" |
| 56 | |
| 57 | |
| 58 | class TemplateChoice(str, Enum): |
| 59 | CORRECT = "correct" |
| 60 | WRONG = "wrong" |
| 61 | PARTIAL = "partial" |
| 62 | |
| 63 | |
| 64 | class LessonType(str, Enum): |
| 65 | ENVIRONMENT = "environment" |
| 66 | FALSE_POSITIVES = "false_positives" |
| 67 | ASSETS = "assets" |
| 68 | THREAT_INTEL = "threat_intel" |
| 69 | ALERTS = "alerts" |
| 70 | |
| 71 | |
| 72 | class Durability(str, Enum): |
| 73 | ONE_OFF = "one_off" |
| 74 | DURABLE = "durable" |
| 75 | |
| 76 | |
| 77 | class PalaceLessonStatus(str, Enum): |
| 78 | PENDING = "pending" |
| 79 | INGESTED = "ingested" |
| 80 | FAILED = "failed" |
| 81 | |
| 82 | |
| 83 | # --- Request schemas --- |
| 84 | |
| 85 | |
| 86 | class CreateJobRequest(BaseModel): |
| 87 | id: str = Field(..., max_length=64, description="Unique job identifier, e.g. copilot-inv-1234-abc") |
| 88 | alert_id: int = Field(..., description="The alert ID from incident_management_alert") |
| 89 | customer_code: str = Field(..., max_length=64, description="Customer code") |
| 90 | triggered_by: TriggeredBy = Field(..., description="How the investigation was triggered") |
| 91 | alert_type: Optional[str] = Field(None, max_length=64, description="Detected alert type, e.g. sysmon_event_1") |
| 92 | template_used: Optional[str] = Field(None, max_length=128, description="Template file used for investigation") |
| 93 | |
| 94 | |
| 95 | class UpdateJobRequest(BaseModel): |
| 96 | status: JobStatus = Field(..., description="New job status") |
| 97 | alert_type: Optional[str] = Field(None, max_length=64) |
| 98 | template_used: Optional[str] = Field(None, max_length=128) |
| 99 | error_message: Optional[str] = Field(None, description="Error message if status is failed") |
| 100 | |
| 101 | |
| 102 | class SubmitReportRequest(BaseModel): |
| 103 | job_id: str = Field(..., max_length=64, description="The job ID this report belongs to") |
| 104 | alert_id: int = Field(..., description="The alert ID") |
| 105 | customer_code: str = Field(..., max_length=64, description="Customer code") |
| 106 | severity_assessment: Optional[SeverityAssessment] = Field(None, description="Severity assessment of the alert") |
| 107 | summary: Optional[str] = Field(None, description="Short summary of findings") |
| 108 | report_markdown: Optional[str] = Field(None, description="Full investigation report in Markdown") |
| 109 | recommended_actions: Optional[str] = Field(None, description="Recommended response actions") |
| 110 | |
| 111 | @field_validator("summary", "report_markdown", "recommended_actions", mode="before") |
| 112 | @classmethod |
| 113 | def strip_control_characters(cls, v): |
| 114 | """Strip control characters that break JSON serialization. |
| 115 | Preserves newline (0x0a), carriage return (0x0d), and tab (0x09). |
| 116 | """ |
| 117 | if v is None: |
| 118 | return v |
| 119 | return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", v) |
| 120 | |
| 121 | @model_validator(mode="after") |
| 122 | def require_report_body(self): |
| 123 | """Reject reports missing any human-readable body field. |
| 124 | |
| 125 | The agent populates these from CLAUDE.md context; when a long-lived |
| 126 | session compacts, the field list can drop out and the agent calls this |
| 127 | with only job_id/alert_id/customer_code. Without this guard the row |
| 128 | persists with NULL bodies and surfaces as an empty report in CoPilot. |
| 129 | Runs after strip_control_characters, so all-control-char values that |
| 130 | collapse to "" are caught here too. |
| 131 | """ |
| 132 | missing = [ |
| 133 | name |
| 134 | for name in ("severity_assessment", "summary", "report_markdown", "recommended_actions") |
| 135 | if not (getattr(self, name) and str(getattr(self, name)).strip()) |
| 136 | ] |
| 137 | if missing: |
| 138 | raise ValueError( |
| 139 | f"Report body fields must be non-empty: {', '.join(missing)}. " "A report missing these persists a blank row in CoPilot.", |
| 140 | ) |
| 141 | return self |
| 142 | |
| 143 | |
| 144 | class SubmitIocRequest(BaseModel): |
| 145 | ioc_value: str = Field(..., max_length=512, description="The IOC value") |
| 146 | ioc_type: IocType = Field(..., description="Type of IOC") |
| 147 | vt_verdict: VtVerdict = Field(default=VtVerdict.UNKNOWN, description="VirusTotal verdict") |
| 148 | vt_score: Optional[str] = Field(None, max_length=32, description="VirusTotal score, e.g. 5/70") |
| 149 | details: Optional[str] = Field(None, description="Additional enrichment details") |
| 150 | |
| 151 | |
| 152 | class SubmitIocsRequest(BaseModel): |
| 153 | report_id: int = Field(..., description="The report ID these IOCs belong to") |
| 154 | alert_id: int = Field(..., description="The alert ID") |
| 155 | customer_code: str = Field(..., max_length=64, description="Customer code") |
| 156 | iocs: List[SubmitIocRequest] = Field(..., description="List of IOCs to submit") |
| 157 | |
| 158 | |
| 159 | # --- Response schemas --- |
| 160 | |
| 161 | |
| 162 | class JobResponse(BaseModel): |
| 163 | id: str |
| 164 | alert_id: int |
| 165 | customer_code: str |
| 166 | status: str |
| 167 | alert_type: Optional[str] = None |
| 168 | triggered_by: str |
| 169 | template_used: Optional[str] = None |
| 170 | created_at: datetime |
| 171 | started_at: Optional[datetime] = None |
| 172 | completed_at: Optional[datetime] = None |
| 173 | error_message: Optional[str] = None |
| 174 | |
| 175 | |
| 176 | class ReportResponse(BaseModel): |
| 177 | id: int |
| 178 | job_id: str |
| 179 | alert_id: int |
| 180 | customer_code: str |
| 181 | severity_assessment: Optional[str] = None |
| 182 | summary: Optional[str] = None |
| 183 | report_markdown: Optional[str] = None |
| 184 | recommended_actions: Optional[str] = None |
| 185 | created_at: datetime |
| 186 | |
| 187 | |
| 188 | class IocResponse(BaseModel): |
| 189 | id: int |
| 190 | report_id: int |
| 191 | alert_id: int |
| 192 | customer_code: str |
| 193 | ioc_value: str |
| 194 | ioc_type: str |
| 195 | vt_verdict: str |
| 196 | vt_score: Optional[str] = None |
| 197 | details: Optional[str] = None |
| 198 | created_at: datetime |
| 199 | |
| 200 | |
| 201 | class CreateJobResponse(BaseModel): |
| 202 | success: bool |
| 203 | message: str |
| 204 | job: Optional[JobResponse] = None |
| 205 | |
| 206 | |
| 207 | class UpdateJobResponse(BaseModel): |
| 208 | success: bool |
| 209 | message: str |
| 210 | job: Optional[JobResponse] = None |
| 211 | |
| 212 | |
| 213 | class SubmitReportResponse(BaseModel): |
| 214 | success: bool |
| 215 | message: str |
| 216 | report: Optional[ReportResponse] = None |
| 217 | |
| 218 | |
| 219 | class SubmitIocsResponse(BaseModel): |
| 220 | success: bool |
| 221 | message: str |
| 222 | iocs_created: int = 0 |
| 223 | iocs: Optional[List[IocResponse]] = None |
| 224 | |
| 225 | |
| 226 | class JobListResponse(BaseModel): |
| 227 | success: bool |
| 228 | message: str |
| 229 | jobs: List[JobResponse] |
| 230 | |
| 231 | |
| 232 | class ReportListResponse(BaseModel): |
| 233 | success: bool |
| 234 | message: str |
| 235 | reports: List[ReportResponse] |
| 236 | |
| 237 | |
| 238 | class IocListResponse(BaseModel): |
| 239 | success: bool |
| 240 | message: str |
| 241 | iocs: List[IocResponse] |
| 242 | |
| 243 | |
| 244 | class AlertWithReportResponse(BaseModel): |
| 245 | alert_id: int |
| 246 | alert_name: str |
| 247 | customer_code: str |
| 248 | status: str |
| 249 | source: str |
| 250 | assigned_to: Optional[str] = None |
| 251 | alert_creation_time: datetime |
| 252 | report: ReportResponse |
| 253 | |
| 254 | |
| 255 | class AlertsWithReportsListResponse(BaseModel): |
| 256 | success: bool |
| 257 | message: str |
| 258 | alerts: List[AlertWithReportResponse] |
| 259 | |
| 260 | |
| 261 | class AlertAnalysisResponse(BaseModel): |
| 262 | success: bool |
| 263 | message: str |
| 264 | job: Optional[JobResponse] = None |
| 265 | report: Optional[ReportResponse] = None |
| 266 | iocs: Optional[List[IocResponse]] = None |
| 267 | |
| 268 | |
| 269 | # --- Review / Palace Lesson / Replay schemas --- |
| 270 | |
| 271 | |
| 272 | class IocVerdictCorrection(BaseModel): |
| 273 | ioc_id: int = Field(..., description="The AiAnalystIoc.id being reviewed") |
| 274 | verdict_correct: bool = Field(..., description="True if the original VT verdict was correct") |
| 275 | note: Optional[str] = Field(None, max_length=2000, description="Optional reviewer note") |
| 276 | |
| 277 | @field_validator("note", mode="before") |
| 278 | @classmethod |
| 279 | def strip_control_characters_note(cls, v): |
| 280 | if v is None: |
| 281 | return v |
| 282 | return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", v) |
| 283 | |
| 284 | |
| 285 | class SubmitReviewRequest(BaseModel): |
| 286 | overall_verdict: Optional[OverallVerdict] = Field(None, description="Overall thumbs up/down") |
| 287 | template_choice: Optional[TemplateChoice] = Field(None, description="Was the selected template correct") |
| 288 | template_used: Optional[str] = Field(None, max_length=128, description="Template filename that ran (mirrored from report)") |
| 289 | rating_instructions: Optional[int] = Field(None, ge=1, le=5, description="Rating 1–5 on instructions quality") |
| 290 | rating_artifacts: Optional[int] = Field(None, ge=1, le=5, description="Rating 1–5 on collected artifacts") |
| 291 | rating_severity: Optional[int] = Field(None, ge=1, le=5, description="Rating 1–5 on severity assessment accuracy") |
| 292 | missing_steps: Optional[str] = Field(None, description="Free-text list of steps the analyst missed") |
| 293 | suggested_edits: Optional[str] = Field(None, description="Free-text suggested prompt / template edits") |
| 294 | ioc_reviews: List[IocVerdictCorrection] = Field(default_factory=list, description="Per-IOC verdict corrections") |
| 295 | |
| 296 | @field_validator("missing_steps", "suggested_edits", mode="before") |
| 297 | @classmethod |
| 298 | def strip_control_characters(cls, v): |
| 299 | if v is None: |
| 300 | return v |
| 301 | return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", v) |
| 302 | |
| 303 | |
| 304 | class IocReviewResponse(BaseModel): |
| 305 | id: int |
| 306 | review_id: int |
| 307 | ioc_id: int |
| 308 | verdict_correct: bool |
| 309 | note: Optional[str] = None |
| 310 | created_at: datetime |
| 311 | |
| 312 | |
| 313 | class ReviewResponse(BaseModel): |
| 314 | id: int |
| 315 | report_id: int |
| 316 | alert_id: int |
| 317 | customer_code: str |
| 318 | reviewer_user_id: int |
| 319 | overall_verdict: Optional[str] = None |
| 320 | template_choice: Optional[str] = None |
| 321 | template_used: Optional[str] = None |
| 322 | rating_instructions: Optional[int] = None |
| 323 | rating_artifacts: Optional[int] = None |
| 324 | rating_severity: Optional[int] = None |
| 325 | missing_steps: Optional[str] = None |
| 326 | suggested_edits: Optional[str] = None |
| 327 | created_at: datetime |
| 328 | updated_at: Optional[datetime] = None |
| 329 | ioc_reviews: List[IocReviewResponse] = Field(default_factory=list) |
| 330 | |
| 331 | |
| 332 | class MyReviewResponse(BaseModel): |
| 333 | """Response for 'fetch my existing review for this report' — used by the UI to |
| 334 | decide whether to show the rubric in create mode or edit-existing mode.""" |
| 335 | |
| 336 | success: bool |
| 337 | message: str |
| 338 | review: Optional[ReviewResponse] = None |
| 339 | |
| 340 | |
| 341 | class SubmitReviewResponse(BaseModel): |
| 342 | success: bool |
| 343 | message: str |
| 344 | review: Optional[ReviewResponse] = None |
| 345 | |
| 346 | |
| 347 | class ReviewListResponse(BaseModel): |
| 348 | success: bool |
| 349 | message: str |
| 350 | reviews: List[ReviewResponse] |
| 351 | |
| 352 | |
| 353 | class QueuePalaceLessonRequest(BaseModel): |
| 354 | customer_code: str = Field(..., max_length=64, description="Customer code this lesson applies to") |
| 355 | lesson_type: LessonType = Field(..., description="MemPalace room / category") |
| 356 | lesson_text: str = Field(..., min_length=1, description="The lesson text to store") |
| 357 | durability: Durability = Field(default=Durability.DURABLE, description="one_off = single-session hint, durable = persistent knowledge") |
| 358 | review_id: Optional[int] = Field(None, description="Optional review.id this lesson was born from") |
| 359 | |
| 360 | @field_validator("lesson_text", mode="before") |
| 361 | @classmethod |
| 362 | def strip_control_characters_lesson(cls, v): |
| 363 | if v is None: |
| 364 | return v |
| 365 | return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", v) |
| 366 | |
| 367 | |
| 368 | class PalaceLessonResponse(BaseModel): |
| 369 | id: int |
| 370 | review_id: Optional[int] = None |
| 371 | customer_code: str |
| 372 | lesson_type: str |
| 373 | lesson_text: str |
| 374 | durability: str |
| 375 | status: str |
| 376 | ingested_at: Optional[datetime] = None |
| 377 | created_at: datetime |
| 378 | |
| 379 | |
| 380 | class QueuePalaceLessonResponse(BaseModel): |
| 381 | success: bool |
| 382 | message: str |
| 383 | lesson: Optional[PalaceLessonResponse] = None |
| 384 | |
| 385 | |
| 386 | class ReplayRequest(BaseModel): |
| 387 | template_override: str = Field( |
| 388 | ..., |
| 389 | max_length=128, |
| 390 | description="Template filename to force for this replay (e.g. sysmon_event_1.txt)", |
| 391 | ) |
| 392 | customer_code: str = Field(..., max_length=64, description="Customer code for the alert") |
| 393 | sender: str = Field(default="copilot-replay", max_length=64, description="Sender identifier for audit") |
| 394 | |
| 395 | @field_validator("template_override") |
| 396 | @classmethod |
| 397 | def validate_template_filename(cls, v): |
| 398 | if not re.match(r"^[a-zA-Z0-9._-]+\.txt$", v): |
| 399 | raise ValueError("template_override must be a filename matching ^[a-zA-Z0-9._-]+\\.txt$") |
| 400 | return v |
| 401 | |
| 402 | |
| 403 | class ReplayResponse(BaseModel): |
| 404 | success: bool |
| 405 | message: str |
| 406 | data: Optional[dict] = None |
| 407 | |
| 408 | |
| 409 | class PalaceSearchHit(BaseModel): |
| 410 | id: Optional[str] = None |
| 411 | room: Optional[str] = None |
| 412 | wing: Optional[str] = None |
| 413 | text: Optional[str] = None |
| 414 | source_file: Optional[str] = None |
| 415 | score: Optional[float] = None |
| 416 | metadata: Optional[dict] = None |
| 417 | |
| 418 | |
| 419 | class PalaceSearchResponse(BaseModel): |
| 420 | success: bool |
| 421 | message: str |
| 422 | lessons: List[PalaceSearchHit] = Field(default_factory=list) |
| 423 | |
| 424 | |
| 425 | # --- Review stats / feedback dashboard --- |
| 426 | |
| 427 | |
| 428 | class ReviewStatsTemplate(BaseModel): |
| 429 | """Per-template slice of review metrics (grouped by template_used).""" |
| 430 | |
| 431 | template_used: Optional[str] = Field(None, description="Template filename, or None for untemplated runs") |
| 432 | total: int = 0 |
| 433 | thumbs_up: int = 0 |
| 434 | thumbs_down: int = 0 |
| 435 | correct: int = 0 |
| 436 | partial: int = 0 |
| 437 | wrong: int = 0 |
| 438 | avg_rating_instructions: Optional[float] = None |
| 439 | avg_rating_artifacts: Optional[float] = None |
| 440 | avg_rating_severity: Optional[float] = None |
| 441 | |
| 442 | |
| 443 | class ReviewStatsIocAccuracy(BaseModel): |
| 444 | """Aggregate IOC verdict accuracy — derived from analyst per-IOC corrections.""" |
| 445 | |
| 446 | total: int = 0 |
| 447 | correct: int = 0 |
| 448 | incorrect: int = 0 |
| 449 | accuracy_pct: Optional[float] = None |
| 450 | |
| 451 | |
| 452 | class ReviewStatsResponse(BaseModel): |
| 453 | success: bool |
| 454 | message: str |
| 455 | customer_code: str |
| 456 | total_reviews: int = 0 |
| 457 | thumbs_up: int = 0 |
| 458 | thumbs_down: int = 0 |
| 459 | thumbs_up_pct: Optional[float] = None |
| 460 | template_choice_correct: int = 0 |
| 461 | template_choice_partial: int = 0 |
| 462 | template_choice_wrong: int = 0 |
| 463 | avg_rating_instructions: Optional[float] = None |
| 464 | avg_rating_artifacts: Optional[float] = None |
| 465 | avg_rating_severity: Optional[float] = None |
| 466 | ioc_accuracy: ReviewStatsIocAccuracy = Field(default_factory=ReviewStatsIocAccuracy) |
| 467 | per_template: List[ReviewStatsTemplate] = Field(default_factory=list) |
| 468 | recent_reviews: List[ReviewResponse] = Field(default_factory=list) |
| 469 | |
| 470 | |
| 471 | # --- Palace consolidation (Step 21.B) --- |
| 472 | |
| 473 | |
| 474 | class PalaceConsolidationLesson(BaseModel): |
| 475 | """A single lesson row, shaped for the consolidation digest UI.""" |
| 476 | |
| 477 | id: int |
| 478 | lesson_type: str |
| 479 | lesson_text: str |
| 480 | durability: str |
| 481 | status: str |
| 482 | drawer_id: Optional[str] = None |
| 483 | created_at: datetime |
| 484 | ingested_at: Optional[datetime] = None |
| 485 | # For one_off lessons only — how many days until the sweeper expires |
| 486 | # this row. Negative numbers mean the sweeper is about to take it on |
| 487 | # the next tick. None for durable rows (no expiry). |
| 488 | days_until_expiry: Optional[int] = None |
| 489 | |
| 490 | |
| 491 | class PalaceConsolidationRoomGroup(BaseModel): |
| 492 | """Per-room slice — lessons grouped by lesson_type.""" |
| 493 | |
| 494 | room: str |
| 495 | total: int |
| 496 | durable: int |
| 497 | one_off: int |
| 498 | lessons: List[PalaceConsolidationLesson] = Field(default_factory=list) |
| 499 | |
| 500 | |
| 501 | class PalaceConsolidationDuplicatePair(BaseModel): |
| 502 | """Near-duplicate candidate flagged for reviewer attention.""" |
| 503 | |
| 504 | room: str |
| 505 | lesson_a_id: int |
| 506 | lesson_b_id: int |
| 507 | lesson_a_text: str |
| 508 | lesson_b_text: str |
| 509 | similarity: float # 0.0 – 1.0, difflib SequenceMatcher ratio |
| 510 | |
| 511 | |
| 512 | class PalaceConsolidationResponse(BaseModel): |
| 513 | """Full digest for a customer — renders inline in a drawer; the |
| 514 | reviewer can also grab the pre-rendered markdown for export.""" |
| 515 | |
| 516 | success: bool |
| 517 | message: str |
| 518 | customer_code: str |
| 519 | generated_at: datetime |
| 520 | # Top-level counts across active (non-expired, non-failed) lessons |
| 521 | total_lessons: int = 0 |
| 522 | total_durable: int = 0 |
| 523 | total_one_off: int = 0 |
| 524 | total_pending: int = 0 |
| 525 | total_ingested: int = 0 |
| 526 | # One-off lessons whose expiry is within SOON_WINDOW_DAYS — reviewer |
| 527 | # may want to promote them to durable before the sweeper deletes them. |
| 528 | upcoming_expirations: List[PalaceConsolidationLesson] = Field(default_factory=list) |
| 529 | rooms: List[PalaceConsolidationRoomGroup] = Field(default_factory=list) |
| 530 | duplicate_candidates: List[PalaceConsolidationDuplicatePair] = Field(default_factory=list) |
| 531 | # Rendered markdown digest — pre-baked so the drawer can offer a |
| 532 | # "copy as markdown" button without client-side templating. |
| 533 | markdown: str = "" |