main
py 721 lines 15.5 KB
Raw
1 from datetime import datetime
2 from enum import Enum
3 from typing import Dict
4 from typing import List
5 from typing import Optional
6
7 from fastapi import HTTPException
8 from pydantic import BaseModel
9 from pydantic import ConfigDict
10 from pydantic import field_validator
11 from pydantic import model_validator
12
13 from app.incidents.models import Alert
14 from app.incidents.models import AlertContext
15 from app.incidents.models import AlertTag
16 from app.incidents.models import AlertToIoC
17 from app.incidents.models import Asset
18 from app.incidents.models import Case
19 from app.incidents.models import CaseAlertLink
20 from app.incidents.models import CaseComment
21 from app.incidents.models import CaseDataStore
22 from app.incidents.models import CaseReportTemplateDataStore
23 from app.incidents.models import Comment
24
25
26 class DeleteAlertsRequest(BaseModel):
27 alert_ids: List[int]
28
29
30 class SocfortressRecommendsWazuhFieldNames(Enum):
31 # ! Windows Events
32 data_win_eventdata_commandLine = "data_win_eventdata_commandLine"
33 data_win_eventdata_parentCommandLine = "data_win_eventdata_parentCommandLine"
34 data_win_eventdata_parentImage = "data_win_eventdata_parentImage"
35 data_win_eventdata_parentUser = "data_win_eventdata_parentUser"
36 data_win_eventdata_image = "data_win_eventdata_image"
37 data_win_eventdata_user = "data_win_eventdata_user"
38 rule_mitre_id = "rule_mitre_id"
39 rule_mitre_tactic = "rule_mitre_tactic"
40 rule_mitre_technique = "rule_mitre_technique"
41 data_win_eventdata_company = "data_win_eventdata_company"
42 data_win_eventdata_hashes = "data_win_eventdata_hashes"
43 data_win_eventdata_currentDirectory = "data_win_eventdata_currentDirectory"
44 data_win_eventdata_originalFileName = "data_win_eventdata_originalFileName"
45 # ! Windows SIGCHECK HITS
46 data_Path = "data_Path"
47 # ! Extra Use for Within CoPilot
48 process_id = "process_id"
49 sha256 = "sha256"
50
51
52 class DeleteAlertsResponse(BaseModel):
53 message: str
54 deleted_alert_ids: List[int]
55 not_deleted_alert_ids: List[int]
56 success: bool
57
58
59 class SocfortressRecommendsWazuhAssetName(Enum):
60 agent_name = "agent_name"
61
62
63 class SocfortressRecommendsWazuhTimeFieldName(Enum):
64 timestamp_utc = "timestamp_utc"
65
66
67 class SocfortressRecommendsWazuhAlertTitleName(Enum):
68 rule_description = "rule_description"
69
70
71 class SocfortressRecommendsWazuhIoCFieldNames(Enum):
72 threat_intel_value = "threat_intel_value"
73
74
75 class SocfortressRecommendsWazuhResponse(BaseModel):
76 field_names: List[str]
77 asset_name: str
78 timefield_name: str
79 alert_title_name: str
80 ioc_field_names: Optional[List[str]] = None
81 source: str
82 success: bool
83 message: str
84
85
86 class AvailableSourcesResponse(BaseModel):
87 source: str
88 success: bool
89 message: str
90
91
92 class AvailableIndicesResponse(BaseModel):
93 indices: List[str]
94 success: bool
95 message: str
96
97
98 class ConfiguredSourcesResponse(BaseModel):
99 sources: List[str]
100 success: bool
101 message: str
102
103
104 class MappingsResponse(BaseModel):
105 available_mappings: List[str]
106 success: bool
107 message: str
108
109
110 class ValidSources(str, Enum):
111 WAZUH = "wazuh"
112
113
114 class AlertStatus(str, Enum):
115 OPEN = "OPEN"
116 CLOSED = "CLOSED"
117 IN_PROGRESS = "IN_PROGRESS"
118
119
120 class UpdateAlertStatus(BaseModel):
121 alert_id: int
122 status: AlertStatus
123
124
125 class UpdateCaseStatus(BaseModel):
126 case_id: int
127 status: AlertStatus
128
129
130 class AlertResponse(BaseModel):
131 alert: Alert
132 success: bool
133 message: str
134
135
136 class CommentResponse(BaseModel):
137 comment: Comment
138 success: bool
139 message: str
140
141
142 class CaseCommentResponse(BaseModel):
143 comment: CaseComment
144 success: bool
145 message: str
146
147
148 class AlertContextResponse(BaseModel):
149 alert_context: AlertContext
150 success: bool
151 message: str
152
153
154 class AssetResponse(BaseModel):
155 asset: Asset
156 success: bool
157 message: str
158
159
160 class AlertTagResponse(BaseModel):
161 alert_tag: AlertTag
162 success: bool
163 message: str
164
165
166 class AlertIocValue(str, Enum):
167 IP = "IP"
168 DOMAIN = "DOMAIN"
169 HASH = "HASH"
170 URL = "URL"
171
172
173 class AlertIoCCreate(BaseModel):
174 alert_id: int
175 ioc_value: str
176 ioc_type: AlertIocValue
177 ioc_description: Optional[str] = None
178
179 @field_validator("ioc_type")
180 @classmethod
181 def validate_ioc_type(cls, v):
182 if v not in AlertIocValue:
183 raise HTTPException(
184 status_code=400,
185 detail=f"Invalid IoC type. Must be one of {', '.join([ioc.value for ioc in AlertIocValue])}",
186 )
187 return v
188
189
190 class AlertIoCResponse(BaseModel):
191 alert_ioc: AlertToIoC
192 success: bool
193 message: str
194
195
196 class CaseResponse(BaseModel):
197 case: Case
198 success: bool
199 message: str
200
201
202 class CaseAlertLinkResponse(BaseModel):
203 case_alert_link: CaseAlertLink
204 success: bool
205 message: str
206
207
208 class CaseAlertUnLinkResponse(BaseModel):
209 success: bool
210 message: str
211 tasks_orphaned: int = 0
212
213
214 class CaseAlertLinksResponse(BaseModel):
215 case_alert_links: List[CaseAlertLink]
216 success: bool
217 message: str
218
219
220 class AvailableUsersResponse(BaseModel):
221 available_users: List[str]
222 success: bool
223 message: str
224
225
226 class FieldAndAssetNames(BaseModel):
227 field_names: List[str]
228 asset_name: str
229 timefield_name: str
230 alert_title_name: str
231 ioc_field_names: Optional[List[str]] = None
232 source: str
233
234
235 class FieldAndAssetNamesResponse(BaseModel):
236 field_names: List[str]
237 asset_name: str
238 timefield_name: str
239 alert_title_name: str
240 ioc_field_names: Optional[List[str]] = None
241 source: str
242 success: bool
243 message: str
244
245
246 class AssignedToAlert(BaseModel):
247 alert_id: int
248 assigned_to: str
249
250
251 class AssignedToCase(BaseModel):
252 case_id: int
253 assigned_to: str
254
255
256 class EscalateAlert(BaseModel):
257 alert_id: int
258 escalated: bool
259
260
261 class EscalateCase(BaseModel):
262 case_id: int
263 escalated: bool
264
265
266 class AlertCreate(BaseModel):
267 alert_name: str
268 alert_description: str
269 status: str
270 alert_creation_time: datetime
271 customer_code: str
272 time_closed: Optional[datetime] = None
273 source: str
274 assigned_to: str
275
276
277 class CommentCreate(BaseModel):
278 alert_id: int
279 comment: str
280 user_name: str
281 created_at: Optional[datetime] = None
282
283
284 class CommentEdit(BaseModel):
285 alert_id: int
286 comment_id: int
287 comment: str
288 user_name: str
289 created_at: datetime
290
291
292 class CaseCommentCreate(BaseModel):
293 case_id: int
294 comment: str
295 user_name: str
296 created_at: Optional[datetime] = None
297
298
299 class CaseCommentEdit(BaseModel):
300 case_id: int
301 comment_id: int
302 comment: str
303 user_name: str
304 created_at: datetime
305
306
307 class AlertContextCreate(BaseModel):
308 source: str
309 context: Dict
310
311
312 class CaseCreate(BaseModel):
313 case_name: str
314 case_description: str
315 case_creation_time: datetime
316 case_status: str
317 assigned_to: Optional[str] = None
318 customer_code: Optional[str] = None
319
320
321 class LinkedCaseCreate(BaseModel):
322 case_name: str
323 case_description: str
324 case_creation_time: str
325 case_status: str
326 assigned_to: Optional[str] = None
327 id: int
328
329 @field_validator("case_creation_time", mode="before")
330 @classmethod
331 def format_case_creation_time(cls, v):
332 if isinstance(v, datetime):
333 return v.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
334 return v
335
336
337 class CaseCreateFromAlert(BaseModel):
338 alert_id: int
339
340
341 class CaseAlertLinkCreate(BaseModel):
342 case_id: int
343 alert_id: int
344
345
346 class CaseAlertLinksCreate(BaseModel):
347 case_id: int
348 alert_ids: List[int]
349
350
351 class CaseAlertUnLink(BaseModel):
352 case_id: int
353 alert_id: int
354
355
356 class AssetCreate(BaseModel):
357 alert_linked: int
358 asset_name: str
359 alert_context_id: int
360 agent_id: Optional[str] = None
361 velociraptor_id: Optional[str] = None
362 customer_code: str
363 index_name: str
364 index_id: str
365
366
367 class AlertTagBase(BaseModel):
368 tag: str
369 id: int
370
371
372 class AlertTagCreate(BaseModel):
373 alert_id: int
374 tag: str
375
376
377 class AlertTagDelete(BaseModel):
378 alert_id: int
379 tag_id: int
380
381
382 class AlertIoCDelete(BaseModel):
383 alert_id: int
384 ioc_id: int
385
386
387 class CommentBase(BaseModel):
388 user_name: str
389 alert_id: int
390 id: int
391 comment: str
392 created_at: str
393
394 @field_validator("created_at", mode="before")
395 @classmethod
396 def format_created_at(cls, v):
397 if isinstance(v, datetime):
398 return v.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
399 return v
400
401
402 class CaseCommentBase(BaseModel):
403 user_name: str
404 case_id: int
405 id: int
406 comment: str
407 created_at: str
408
409 @field_validator("created_at", mode="before")
410 @classmethod
411 def format_created_at(cls, v):
412 if isinstance(v, datetime):
413 return v.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
414 return v
415
416
417 class AssetBase(BaseModel):
418 asset_name: str
419 agent_id: Optional[str] = None
420 customer_code: str
421 index_id: str
422 alert_linked: int
423 id: int
424 alert_context_id: int
425 velociraptor_id: Optional[str] = None
426 index_name: str
427
428
429 class IoCBase(BaseModel):
430 value: str
431 type: AlertIocValue
432 description: Optional[str] = None
433 id: int
434
435
436 class AlertOut(BaseModel):
437 id: int
438 alert_creation_time: str
439 time_closed: Optional[str] = None
440 alert_name: str
441 alert_description: str
442 status: str
443 customer_code: str
444 source: str
445 assigned_to: Optional[str] = None
446 escalated: bool = False
447 comments: List[CommentBase] = []
448 assets: List[AssetBase] = []
449 tags: List[AlertTagBase] = []
450 linked_cases: List[LinkedCaseCreate] = []
451 iocs: List[IoCBase] = []
452
453 @field_validator("alert_creation_time", "time_closed", mode="before")
454 @classmethod
455 def format_datetime(cls, v):
456 if isinstance(v, datetime):
457 return v.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
458 return v
459
460
461 class AlertOutResponse(BaseModel):
462 alerts: List[AlertOut]
463 total: Optional[int] = None
464 open: Optional[int] = None
465 in_progress: Optional[int] = None
466 closed: Optional[int] = None
467 total_filtered: Optional[int] = None
468 success: bool
469 message: str
470
471
472 class CaseOut(BaseModel):
473 id: int
474 case_name: str
475 case_description: str
476 assigned_to: Optional[str] = None
477 alerts: Optional[List[AlertOut]] = []
478 case_status: Optional[str] = None
479 case_creation_time: Optional[str] = None
480 customer_code: Optional[str] = None
481 notification_invoked_number: Optional[int] = 0
482 escalated: bool = False
483 comments: List[CaseCommentBase] = []
484
485 @field_validator("case_creation_time", mode="before")
486 @classmethod
487 def format_case_creation_time(cls, v):
488 if isinstance(v, datetime):
489 return v.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
490 return v
491
492
493 class CaseOutResponse(BaseModel):
494 cases: List[CaseOut]
495 total: Optional[int] = None
496 open: Optional[int] = None
497 in_progress: Optional[int] = None
498 closed: Optional[int] = None
499 success: bool
500 message: str
501
502
503 class CaseNotificationCreate(BaseModel):
504 case_id: int
505
506
507 class CaseNotificationResponse(BaseModel):
508 success: bool
509 message: str
510
511
512 class Notification(BaseModel):
513 id: int
514 customer_code: str
515 shuffle_workflow_id: str
516 enabled: bool
517
518 model_config = ConfigDict(from_attributes=True)
519
520
521 class NotificationResponse(BaseModel):
522 notifications: Optional[List[Notification]] = []
523 success: bool
524 message: str
525
526
527 class PutNotification(BaseModel):
528 customer_code: str
529 shuffle_workflow_id: str
530 enabled: bool
531
532
533 class AITrigger(BaseModel):
534 id: int
535 customer_code: str
536 enabled: bool
537
538 model_config = ConfigDict(from_attributes=True)
539
540
541 class AITriggerResponse(BaseModel):
542 ai_triggers: Optional[List[AITrigger]] = []
543 success: bool
544 message: str
545
546
547 class PutAITrigger(BaseModel):
548 customer_code: str
549 enabled: bool
550
551
552 class CaseDataStoreResponse(BaseModel):
553 case_data_store: CaseDataStore
554 success: bool
555 message: str
556
557
558 class ListCaseDataStoreResponse(BaseModel):
559 case_data_store: List[CaseDataStore]
560 success: bool
561 message: str
562
563
564 class CaseReportTemplateDataStoreResponse(BaseModel):
565 case_report_template_data_store: CaseReportTemplateDataStore
566 success: bool
567 message: str
568
569
570 class CaseDownloadDocxRequest(BaseModel):
571 case_id: int
572 template_name: str
573 file_name: Optional[str] = "case_report.docx"
574
575 @field_validator("file_name", mode="before")
576 @classmethod
577 def ensure_docx_extension(cls, v):
578 if v and not v.endswith(".docx"):
579 return f"{v}.docx"
580 return v
581
582
583 class CaseReportTemplateDataStoreListResponse(BaseModel):
584 case_report_template_data_store: List[str]
585 success: bool
586 message: str
587
588
589 class DefaultReportTemplateFileNames(Enum):
590 CASE_REPORT_JINJA_TEMPLATE = "case_report_jinja_template.docx"
591
592
593 class AlertFilterOptionsResponse(BaseModel):
594 sources: List[str]
595 assets: List[str]
596 tags: List[str]
597 statuses: List[str] = [s.value for s in AlertStatus]
598 success: bool
599 message: str
600
601
602 class CaseFilterOptionsResponse(BaseModel):
603 statuses: List[str]
604 assigned_to: List[str]
605 success: bool
606 message: str
607
608
609 # ============================================
610 # Tag Access RBAC Schemas
611 # ============================================
612
613
614 class AlertTagItem(BaseModel):
615 """Single tag item for responses."""
616
617 id: int
618 tag: str
619
620
621 class TagAccessCreate(BaseModel):
622 """Base schema for creating tag access."""
623
624 tag_ids: List[int]
625
626
627 class UserTagAccessCreate(TagAccessCreate):
628 """Assign tags to a user."""
629
630 user_id: int
631
632
633 class RoleTagAccessCreate(TagAccessCreate):
634 """Assign tags to a role."""
635
636 role_id: int
637
638
639 class UserTagAccessResponse(BaseModel):
640 """Response for user tag access operations."""
641
642 user_id: int
643 username: str
644 accessible_tags: List[AlertTagItem]
645 success: bool
646 message: str
647
648
649 class RoleTagAccessResponse(BaseModel):
650 """Response for role tag access operations."""
651
652 role_id: int
653 role_name: str
654 accessible_tags: List[AlertTagItem]
655 success: bool
656 message: str
657
658
659 class UntaggedAlertBehavior(str, Enum):
660 """Options for handling untagged alerts when tag RBAC is enabled."""
661
662 VISIBLE_TO_ALL = "visible_to_all"
663 ADMIN_ONLY = "admin_only"
664 DEFAULT_TAG = "default_tag"
665
666
667 class TagAccessSettingsUpdate(BaseModel):
668 """Update tag access settings."""
669
670 enabled: bool
671 untagged_alert_behavior: UntaggedAlertBehavior = UntaggedAlertBehavior.VISIBLE_TO_ALL
672 default_tag_id: Optional[int] = None
673
674 @model_validator(mode="after")
675 def validate_default_tag(self):
676 if self.untagged_alert_behavior == UntaggedAlertBehavior.DEFAULT_TAG and self.default_tag_id is None:
677 raise HTTPException(
678 status_code=400,
679 detail="default_tag_id is required when untagged_alert_behavior is 'default_tag'",
680 )
681 return self
682
683
684 class TagAccessSettingsItem(BaseModel):
685 """Single tag access settings item."""
686
687 enabled: bool
688 untagged_alert_behavior: str
689 default_tag_id: Optional[int] = None
690 default_tag_name: Optional[str] = None
691
692
693 class TagAccessSettingsResponse(BaseModel):
694 """Response for tag access settings."""
695
696 settings: TagAccessSettingsItem
697 success: bool
698 message: str
699
700
701 class UserEffectiveAccessResponse(BaseModel):
702 """Shows effective access for a user (combines role + user-specific access)."""
703
704 user_id: int
705 username: str
706 role_id: Optional[int] = None
707 role_name: Optional[str] = None
708 accessible_customers: List[str]
709 accessible_tags: List[AlertTagItem]
710 is_tag_unrestricted: bool
711 tag_rbac_enabled: bool
712 success: bool
713 message: str
714
715
716 class AllTagsResponse(BaseModel):
717 """Response for listing all available tags."""
718
719 tags: List[AlertTagItem]
720 success: bool
721 message: str