main
py 390 lines 13.6 KB
Raw
1 """
2 Pydantic schemas for the notification routing module.
3
4 The wire-level enums (NotificationTrigger, NotificationChannel,
5 NotificationSeverity) mirror the v1 string set the database column
6 accepts. The DB columns themselves are plain strings so adding a new
7 trigger or channel later is a data-only change — these enums exist
8 purely for input validation at the API boundary.
9 """
10
11 from __future__ import annotations
12
13 from datetime import datetime
14 from enum import Enum
15 from typing import List
16 from typing import Optional
17
18 from pydantic import BaseModel
19 from pydantic import ConfigDict
20 from pydantic import Field
21 from pydantic import field_validator
22 from pydantic import model_validator
23
24 # ---------------------------------------------------------------------------
25 # Enums (input validation only — DB stores strings)
26 # ---------------------------------------------------------------------------
27
28
29 class NotificationTrigger(str, Enum):
30 """What kind of event caused this dispatch.
31
32 Currently a single value — `investigation_complete` covers every
33 Talon-driven dispatch (one per investigation that reaches the
34 write-back step). Severity-based filtering lives entirely in the
35 route's `min_severity` field, not here, so the trigger is purely
36 an event-type dimension that grows when we add new dispatch
37 sources (analyst-review hooks, scheduled-sweep findings,
38 IOC-enrichment alerts, etc.).
39 """
40
41 INVESTIGATION_COMPLETE = "investigation_complete"
42
43
44 class NotificationChannel(str, Enum):
45 """Delivery channel set.
46
47 `shuffle` proxies to Shuffle's hosted MCP — each customer points at
48 their own Shuffle Org via `customer_shuffle_integration`, and Shuffle
49 handles the OAuth-authenticated downstream app (Slack workspace,
50 Outlook tenant, Teams, Gmail, SendGrid, etc.). Routes referencing
51 `shuffle` MUST populate the `shuffle_integration_id` + `shuffle_app_id`
52 columns. Email, chat, ticketing, and the rest of the catalog all
53 flow through this single channel — there's no separate direct SMTP
54 path because Shuffle's email apps cover that surface.
55 """
56
57 SHUFFLE = "shuffle"
58
59
60 class NotificationSeverity(str, Enum):
61 """Severity tiers, ordered. Mirrors AiAnalystReport.severity_assessment.
62
63 The dispatch service treats `min_severity` inclusively — a route
64 with `min_severity="High"` fires on Critical and High but not Medium.
65 """
66
67 CRITICAL = "Critical"
68 HIGH = "High"
69 MEDIUM = "Medium"
70 LOW = "Low"
71 INFORMATIONAL = "Informational"
72
73
74 class DispatchStatus(str, Enum):
75 """Result classes for notification_dispatch_log.status."""
76
77 SENT = "sent"
78 FAILED = "failed"
79 SKIPPED = "skipped"
80
81
82 # Severity ordering for `min_severity` filtering. Index = priority,
83 # higher = more severe. Used by the dispatch service to gate routes.
84 SEVERITY_ORDER: List[str] = [
85 NotificationSeverity.INFORMATIONAL.value,
86 NotificationSeverity.LOW.value,
87 NotificationSeverity.MEDIUM.value,
88 NotificationSeverity.HIGH.value,
89 NotificationSeverity.CRITICAL.value,
90 ]
91
92
93 # ---------------------------------------------------------------------------
94 # Routes — request/response shapes
95 # ---------------------------------------------------------------------------
96
97
98 class NotificationRouteBase(BaseModel):
99 name: str = Field(..., min_length=1, max_length=128, description="Human label for the rule (e.g. 'SOC team Slack #alerts').")
100 trigger: NotificationTrigger
101 channel: NotificationChannel
102
103 @field_validator("trigger", mode="before")
104 @classmethod
105 def _coerce_legacy_trigger(cls, v):
106 """Coerce legacy `severity_critical_or_high` rows on read.
107
108 Older versions of this schema treated trigger as a severity
109 filter; routes saved against that schema have a stale value
110 the new enum no longer accepts. Pydantic validates BEFORE the
111 enum check when `pre=True`, so we rewrite the legacy value to
112 the new event-type value here. The dispatch loop has the same
113 backward-compat in `_trigger_applies` for the route-side
114 comparison; this is the read-API equivalent.
115 """
116 if v == "severity_critical_or_high":
117 return NotificationTrigger.INVESTIGATION_COMPLETE.value
118 return v
119
120 # For SMTP: comma-separated recipient emails. For Shuffle: free-form
121 # destination hint (e.g. '#soc-alerts', 'ir@corp.com') that gets
122 # injected into Shuffle's natural-language input — Shuffle's app
123 # agent figures out how to route it within the authenticated app.
124 destination: str = Field(
125 ...,
126 min_length=1,
127 description="Destination hint for the Shuffle app (channel name, email address, handle).",
128 )
129 min_severity: NotificationSeverity = NotificationSeverity.MEDIUM
130 format_template: Optional[str] = Field(
131 default=None,
132 description="Optional Jinja override for the message body. Leave empty to use the channel default.",
133 )
134 enabled: bool = True
135
136 # Phase 2: Shuffle routing target. Required when channel='shuffle'.
137 # The integration row scopes the dispatch to a specific customer
138 # Shuffle org; the app id + name describe which app within that
139 # org receives the natural-language input.
140 shuffle_integration_id: Optional[int] = Field(
141 default=None,
142 description="ID of the customer_shuffle_integration row (required when channel='shuffle').",
143 )
144 shuffle_app_id: Optional[str] = Field(default=None, description="Shuffle app UUID (required when channel='shuffle').")
145 shuffle_app_name: Optional[str] = Field(
146 default=None,
147 description="Human-readable Shuffle app name cached for the UI list (e.g. 'Slack').",
148 )
149
150 @field_validator("destination")
151 @classmethod
152 def _strip_destination(cls, v: str) -> str:
153 return v.strip()
154
155 @model_validator(mode="after")
156 def _shuffle_fields_required(self):
157 if self.channel == NotificationChannel.SHUFFLE:
158 if not self.shuffle_integration_id:
159 raise ValueError("shuffle_integration_id is required when channel='shuffle'")
160 if not self.shuffle_app_id:
161 raise ValueError("shuffle_app_id is required when channel='shuffle'")
162 return self
163
164
165 class NotificationRouteCreate(NotificationRouteBase):
166 """Body for POST /customers/{code}/notification_routes."""
167
168
169 class NotificationRouteUpdate(BaseModel):
170 """Body for PATCH — every field optional. Mirrors the editable subset
171 of NotificationRouteBase."""
172
173 name: Optional[str] = Field(default=None, min_length=1, max_length=128)
174 trigger: Optional[NotificationTrigger] = None
175 channel: Optional[NotificationChannel] = None
176 destination: Optional[str] = Field(default=None, min_length=1)
177 min_severity: Optional[NotificationSeverity] = None
178 format_template: Optional[str] = None
179 enabled: Optional[bool] = None
180 # Shuffle target — included on PATCH so admins can re-point a route
181 # at a different integration / app without recreating it.
182 shuffle_integration_id: Optional[int] = None
183 shuffle_app_id: Optional[str] = None
184 shuffle_app_name: Optional[str] = None
185
186
187 class NotificationRouteRead(NotificationRouteBase):
188 id: int
189 customer_code: str
190 last_dispatched_at: Optional[datetime] = None
191 dispatch_count: int = 0
192 created_by: Optional[str] = None
193 created_at: datetime
194 updated_at: Optional[datetime] = None
195 model_config = ConfigDict(from_attributes=True)
196
197
198 # ---------------------------------------------------------------------------
199 # Shuffle integrations (Phase 2)
200 # ---------------------------------------------------------------------------
201
202
203 class ShuffleIntegrationBase(BaseModel):
204 display_name: str = Field(..., min_length=1, max_length=128, description="Human label, e.g. 'Acme Production Shuffle'.")
205 shuffle_org_id: str = Field(
206 ...,
207 min_length=1,
208 max_length=64,
209 description="The customer's Shuffle Org-Id. Sent as the Org-Id header on each dispatch.",
210 )
211 enabled: bool = True
212
213 @field_validator("shuffle_org_id")
214 @classmethod
215 def _strip_org(cls, v: str) -> str:
216 return v.strip()
217
218
219 class ShuffleIntegrationCreate(ShuffleIntegrationBase):
220 """Body for POST /customers/{code}/shuffle_integrations."""
221
222
223 class ShuffleIntegrationUpdate(BaseModel):
224 """Body for PATCH — every field optional."""
225
226 display_name: Optional[str] = Field(default=None, min_length=1, max_length=128)
227 shuffle_org_id: Optional[str] = Field(default=None, min_length=1, max_length=64)
228 enabled: Optional[bool] = None
229
230
231 class ShuffleIntegrationRead(ShuffleIntegrationBase):
232 id: int
233 customer_code: str
234 last_used_at: Optional[datetime] = None
235 created_by: Optional[str] = None
236 created_at: datetime
237 updated_at: Optional[datetime] = None
238 model_config = ConfigDict(from_attributes=True)
239
240
241 class ShuffleIntegrationListResponse(BaseModel):
242 success: bool = True
243 message: str = "Integrations retrieved"
244 integrations: List[ShuffleIntegrationRead]
245
246
247 class ShuffleIntegrationResponse(BaseModel):
248 success: bool = True
249 message: str = "Integration saved"
250 integration: ShuffleIntegrationRead
251
252
253 class ShuffleApp(BaseModel):
254 """One Shuffle app in the catalog the customer's org has access to.
255
256 Used to populate the route form's app picker. We forward the minimal
257 subset Shuffle returns — enough for the UI to render a recognizable
258 list and for the form to record the (id, name) pair on submit.
259 """
260
261 id: str
262 name: str
263 description: Optional[str] = None
264 large_image: Optional[str] = None
265
266
267 class ShuffleAppListResponse(BaseModel):
268 success: bool = True
269 message: str = "Apps retrieved"
270 apps: List[ShuffleApp]
271
272
273 class ShuffleVerifyResponse(BaseModel):
274 success: bool = True
275 message: str
276 org_id: str
277 app_count: Optional[int] = None
278 error: Optional[str] = None
279
280
281 class ShuffleOrg(BaseModel):
282 """One Shuffle org visible to the deployment's admin Bearer.
283
284 Used to populate the integration form's org-picker dropdown so
285 admins don't have to paste UUIDs. Forwards only the fields the UI
286 needs — Shuffle's full org payload carries a lot of internal state
287 (users, billing, region, sync_config) we don't want leaking
288 through. `creator_org` is empty/falsy on top-level orgs and set to
289 the parent's UUID on sub-orgs, so the UI can label sub-orgs
290 distinctly without an extra round-trip.
291 """
292
293 id: str
294 name: str
295 description: Optional[str] = None
296 role: Optional[str] = None
297 creator_org: Optional[str] = None
298
299
300 class ShuffleOrgListResponse(BaseModel):
301 success: bool = True
302 message: str = "Orgs retrieved"
303 orgs: List[ShuffleOrg]
304
305
306 class NotificationRouteListResponse(BaseModel):
307 success: bool = True
308 message: str = "Routes retrieved"
309 routes: List[NotificationRouteRead]
310
311
312 class NotificationRouteResponse(BaseModel):
313 success: bool = True
314 message: str = "Route saved"
315 route: NotificationRouteRead
316
317
318 # ---------------------------------------------------------------------------
319 # Dispatch log — read-only audit shapes
320 # ---------------------------------------------------------------------------
321
322
323 class DispatchLogRead(BaseModel):
324 id: int
325 customer_code: str
326 alert_id: int
327 route_id: int
328 trigger: str
329 dispatched_at: datetime
330 status: DispatchStatus
331 error_message: Optional[str] = None
332 latency_ms: Optional[int] = None
333 payload_preview: Optional[str] = None
334 shuffle_execution_id: Optional[str] = None
335 model_config = ConfigDict(from_attributes=True)
336
337
338 class DispatchLogListResponse(BaseModel):
339 success: bool = True
340 message: str = "Dispatch log retrieved"
341 entries: List[DispatchLogRead]
342
343
344 # ---------------------------------------------------------------------------
345 # Dispatch endpoint — the one Talon calls
346 # ---------------------------------------------------------------------------
347
348
349 class DispatchRequest(BaseModel):
350 """Body for POST /notifications/dispatch — what Talon sends after
351 completing an investigation. Carries the minimum the dispatch
352 service needs to (a) decide which routes match and (b) format the
353 message body."""
354
355 customer_code: str = Field(..., description="The alert's customer_code — scopes the route lookup.")
356 alert_id: int = Field(..., description="The alert this investigation was for. Used as the idempotency key.")
357 trigger: NotificationTrigger = Field(
358 ...,
359 description="Which trigger Talon thinks applies. The service still re-validates it against the alert's severity.",
360 )
361 severity_assessment: NotificationSeverity = Field(
362 ...,
363 description="The report's assessed severity — used for `min_severity` filtering.",
364 )
365 summary: str = Field(..., description="One-paragraph human-readable summary. Renders into the default template.")
366 report_url: Optional[str] = Field(default=None, description="Deep link back to the full report in CoPilot.")
367 alert_name: Optional[str] = Field(default=None, description="Original alert title for context in the message.")
368
369
370 class DispatchOutcome(BaseModel):
371 route_id: int
372 route_name: str
373 channel: str
374 status: DispatchStatus
375 error_message: Optional[str] = None
376 latency_ms: Optional[int] = None
377 # Shuffle's POST /apps/{id}/mcp returns this on a successful kickoff.
378 # Surfaced in the response so the calling agent (Talon) can include
379 # it in its analyst summary if the dispatch went through Shuffle.
380 shuffle_execution_id: Optional[str] = None
381
382
383 class DispatchResponse(BaseModel):
384 success: bool = True
385 message: str = "Dispatch complete"
386 routes_matched: int
387 dispatched: int
388 skipped: int
389 failed: int
390 outcomes: List[DispatchOutcome]