main
py 364 lines 12.3 KB
Raw
1 """
2 Pydantic schemas for case templates, template tasks, case tasks, and the
3 case event timeline.
4
5 These schemas back the Phase 2+ routes (template CRUD, task lifecycle,
6 timeline GET). The underlying SQLModel rows live in ``app.incidents.models``:
7 ``CaseTemplate``, ``CaseTemplateTask``, ``CaseTask``, ``CaseEvent``.
8 """
9
10 from datetime import datetime
11 from enum import Enum
12 from typing import Any
13 from typing import Dict
14 from typing import List
15 from typing import Optional
16
17 from pydantic import BaseModel
18 from pydantic import ConfigDict
19 from pydantic import Field
20 from pydantic import field_validator
21
22 # ---------------------------------------------------------------------------
23 # Enums
24 # ---------------------------------------------------------------------------
25
26
27 class CaseTaskStatus(str, Enum):
28 """Lifecycle status of a CaseTask."""
29
30 TODO = "TODO"
31 DONE = "DONE"
32 NOT_NECESSARY = "NOT_NECESSARY"
33
34
35 class CaseEventType(str, Enum):
36 """Allowed CaseEvent.event_type values. Kept in sync with the audit hooks
37 added in Phase 4. Stored as a string column so unknown values don't fail
38 reads, but writes should funnel through this enum."""
39
40 CASE_CREATED = "case_created"
41 CASE_STATUS_CHANGED = "case_status_changed"
42 CASE_ASSIGNED = "case_assigned"
43 CASE_ESCALATED = "case_escalated"
44 ALERT_LINKED = "alert_linked"
45 ALERT_UNLINKED = "alert_unlinked"
46 COMMENT_ADDED = "comment_added"
47 TEMPLATE_APPLIED = "template_applied"
48 TASK_ADDED = "task_added"
49 TASK_STATUS_CHANGED = "task_status_changed"
50 TASK_COMMENTED = "task_commented"
51
52
53 # ---------------------------------------------------------------------------
54 # CaseTemplateTask (template-side definition rows)
55 # ---------------------------------------------------------------------------
56
57
58 class CaseTemplateTaskCreate(BaseModel):
59 """Payload for adding a task to a template."""
60
61 title: str = Field(..., max_length=500)
62 description: Optional[str] = None
63 guidelines: Optional[str] = Field(None, description="Best practices / step-by-step guidance for the analyst")
64 mandatory: bool = False
65 order_index: int = Field(0, ge=0)
66
67
68 class CaseTemplateTaskUpdate(BaseModel):
69 """Partial update payload for a template task."""
70
71 title: Optional[str] = Field(None, max_length=500)
72 description: Optional[str] = None
73 guidelines: Optional[str] = None
74 mandatory: Optional[bool] = None
75 order_index: Optional[int] = Field(None, ge=0)
76
77
78 class CaseTemplateTaskResponse(BaseModel):
79 id: int
80 template_id: int
81 title: str
82 description: Optional[str] = None
83 guidelines: Optional[str] = None
84 mandatory: bool
85 order_index: int
86 model_config = ConfigDict(from_attributes=True)
87
88
89 # ---------------------------------------------------------------------------
90 # CaseTemplate
91 # ---------------------------------------------------------------------------
92
93
94 class CaseTemplateCreate(BaseModel):
95 """Payload for creating a new case template (admin/analyst only)."""
96
97 name: str = Field(..., max_length=255)
98 description: Optional[str] = None
99 customer_code: Optional[str] = Field(
100 None,
101 max_length=50,
102 description="Customer this template applies to. Omit for a global template.",
103 )
104 source: Optional[str] = Field(
105 None,
106 max_length=50,
107 description="Alert source this template applies to (e.g., wazuh, velociraptor). Omit for any source.",
108 )
109 is_default: bool = Field(
110 False,
111 description="If true, this template is the fallback within its (customer_code, source) scope.",
112 )
113 match_field: Optional[str] = Field(
114 None,
115 max_length=255,
116 description=(
117 "Optional conditional auto-apply field name on the originating Wazuh document "
118 "(e.g., 'data_win_system_eventID'). Both match_field and match_value must be set "
119 "together — providing one without the other is rejected on create/update."
120 ),
121 )
122 match_value: Optional[str] = Field(
123 None,
124 description="Equality value compared against document[match_field] at auto-apply time.",
125 )
126 tasks: List[CaseTemplateTaskCreate] = Field(
127 default_factory=list,
128 description="Initial task list. More can be added later via the task endpoints.",
129 )
130
131
132 class CaseTemplateUpdate(BaseModel):
133 """Partial update payload for template metadata. Tasks are managed via
134 their own endpoints, not through this schema."""
135
136 name: Optional[str] = Field(None, max_length=255)
137 description: Optional[str] = None
138 customer_code: Optional[str] = Field(None, max_length=50)
139 source: Optional[str] = Field(None, max_length=50)
140 is_default: Optional[bool] = None
141 # Explicit None clears the field (when present in the request); omitted field
142 # leaves the existing value alone. Handled in the service via __fields_set__.
143 match_field: Optional[str] = Field(None, max_length=255)
144 match_value: Optional[str] = None
145
146
147 class CaseTemplateResponse(BaseModel):
148 id: int
149 name: str
150 description: Optional[str] = None
151 customer_code: Optional[str] = None
152 source: Optional[str] = None
153 is_default: bool
154 match_field: Optional[str] = None
155 match_value: Optional[str] = None
156 created_by: str
157 created_at: datetime
158 updated_at: datetime
159 tasks: List[CaseTemplateTaskResponse] = Field(default_factory=list)
160 model_config = ConfigDict(from_attributes=True)
161
162
163 class CaseTemplateListResponse(BaseModel):
164 templates: List[CaseTemplateResponse] = Field(default_factory=list)
165 success: bool
166 message: str
167
168
169 class CaseTemplateOperationResponse(BaseModel):
170 template: Optional[CaseTemplateResponse] = None
171 success: bool
172 message: str
173
174
175 class CaseTemplateTaskOperationResponse(BaseModel):
176 task: Optional[CaseTemplateTaskResponse] = None
177 success: bool
178 message: str
179
180
181 # ---------------------------------------------------------------------------
182 # CaseTask (case-side instance rows)
183 # ---------------------------------------------------------------------------
184
185
186 class CaseTaskCreate(BaseModel):
187 """Payload for adding a custom task to an existing case (analyst-driven)."""
188
189 title: str = Field(..., max_length=500)
190 description: Optional[str] = None
191 guidelines: Optional[str] = None
192 mandatory: bool = False
193 order_index: int = Field(0, ge=0)
194 alert_id: Optional[int] = Field(
195 None,
196 description=(
197 "Optional originating alert. When set, the alert must be linked to the case "
198 "or the request is rejected. Omit (or pass null) to create a case-wide / general task."
199 ),
200 )
201
202
203 class CaseTaskUpdate(BaseModel):
204 """
205 Partial update payload for a case task.
206
207 Status transitions to NOT_NECESSARY are rejected at the service layer
208 when the task is mandatory. ``evidence_comment`` is intended for free-form
209 notes / log snippets / command output captured alongside the status change.
210 """
211
212 status: Optional[CaseTaskStatus] = None
213 evidence_comment: Optional[str] = None
214
215 @field_validator("status")
216 @classmethod
217 def _status_must_be_known(cls, v: Optional[CaseTaskStatus]) -> Optional[CaseTaskStatus]:
218 # Pydantic already enforces enum membership; this guard is for clarity
219 # and to catch any future string-coercion shenanigans.
220 if v is not None and v not in CaseTaskStatus:
221 raise ValueError(f"Unknown task status: {v}")
222 return v
223
224
225 class CaseTaskResponse(BaseModel):
226 id: int
227 case_id: int
228 alert_id: Optional[int] = None
229 template_task_id: Optional[int] = None
230 title: str
231 description: Optional[str] = None
232 guidelines: Optional[str] = None
233 mandatory: bool
234 order_index: int
235 status: CaseTaskStatus
236 evidence_comment: Optional[str] = None
237 completed_by: Optional[str] = None
238 completed_at: Optional[datetime] = None
239 created_by: str
240 created_at: datetime
241 updated_at: datetime
242 model_config = ConfigDict(from_attributes=True)
243
244
245 class CaseTaskListResponse(BaseModel):
246 tasks: List[CaseTaskResponse] = Field(default_factory=list)
247 success: bool
248 message: str
249
250
251 class CaseTaskOperationResponse(BaseModel):
252 task: Optional[CaseTaskResponse] = None
253 success: bool
254 message: str
255
256
257 # ---------------------------------------------------------------------------
258 # Soft-warning payload returned when an analyst tries to close a case with
259 # incomplete mandatory tasks. The route returns this object with HTTP 200 and
260 # the case is NOT closed; the caller re-submits with ?force=true to confirm.
261 # ---------------------------------------------------------------------------
262
263
264 class CaseCloseWarningResponse(BaseModel):
265 """
266 Soft-warning response when closing a case with incomplete mandatory tasks.
267
268 The case is NOT closed when this response is returned. Re-submit the close
269 request with ``force=true`` to override the warning.
270 """
271
272 success: bool = False
273 requires_confirmation: bool = True
274 message: str
275 incomplete_mandatory_tasks: List[CaseTaskResponse] = Field(default_factory=list)
276
277
278 # ---------------------------------------------------------------------------
279 # CaseEvent / Timeline
280 # ---------------------------------------------------------------------------
281
282
283 class CaseEventResponse(BaseModel):
284 id: int
285 case_id: int
286 event_type: str
287 actor: str
288 timestamp: datetime
289 payload: Optional[Dict[str, Any]] = None
290 model_config = ConfigDict(from_attributes=True)
291
292
293 class CaseTimelineResponse(BaseModel):
294 case_id: int
295 events: List[CaseEventResponse] = Field(default_factory=list)
296 success: bool
297 message: str
298
299
300 # ---------------------------------------------------------------------------
301 # Case Template Library — read-only catalog of playbooks pulled from
302 # https://github.com/socfortress/CoPilot-Case-Templates.
303 #
304 # These models mirror the YAML schema documented in that repo's ``SCHEMA.md``.
305 # A LibraryEntry is *not* a CaseTemplate row — it's the YAML view of a
306 # playbook. Admins import an entry via ``POST /library/{key}/import`` which
307 # in turn calls the existing ``create_template`` service, materialising a
308 # normal CaseTemplate (+ task) row set.
309 # ---------------------------------------------------------------------------
310
311
312 class CaseTemplateLibraryTask(BaseModel):
313 """One task in a library entry. Mirrors ``CaseTemplateTaskCreate`` plus
314 explicit ``order_index`` (always present after the loader normalises)."""
315
316 title: str = Field(..., max_length=500)
317 description: Optional[str] = None
318 guidelines: Optional[str] = None
319 mandatory: bool = False
320 order_index: int = Field(..., ge=0)
321
322
323 class CaseTemplateLibraryEntry(BaseModel):
324 """A single playbook surfaced from the Library repo. Display-only on its
325 own; becomes a CaseTemplate row when imported."""
326
327 key: str = Field(..., description="Stable identifier from the YAML; used for collision detection on import")
328 name: str = Field(..., max_length=255)
329 description: Optional[str] = None
330 source: Optional[str] = Field(None, max_length=50)
331 match_field: Optional[str] = Field(
332 None,
333 max_length=255,
334 description="Optional conditional auto-apply field (from the YAML ``match.field``). Persisted on import.",
335 )
336 match_value: Optional[str] = Field(
337 None,
338 description="Optional conditional auto-apply value (from the YAML ``match.value``). Persisted on import.",
339 )
340 tags: Dict[str, Any] = Field(default_factory=dict, description="Library-only metadata; not persisted into the DB")
341 tasks: List[CaseTemplateLibraryTask] = Field(default_factory=list)
342 file_path: Optional[str] = Field(None, description="Path of the source YAML within the repo (for display only)")
343
344
345 class CaseTemplateLibraryListResponse(BaseModel):
346 entries: List[CaseTemplateLibraryEntry] = Field(default_factory=list)
347 invalid_paths: List[str] = Field(
348 default_factory=list,
349 description="Repo paths that failed validation during the last refresh; surfaced so admins can fix upstream YAML.",
350 )
351 last_refresh: Optional[datetime] = Field(
352 None,
353 description="When the cache was last refreshed; null if the cache hasn't loaded yet.",
354 )
355 success: bool
356 message: str
357
358
359 class CaseTemplateLibraryRefreshResponse(BaseModel):
360 loaded: int
361 invalid_paths: List[str] = Field(default_factory=list)
362 last_refresh: Optional[datetime] = None
363 success: bool
364 message: str