main
py 906 lines 43.4 KB
Raw
1 from datetime import datetime
2 from typing import List
3 from typing import Optional
4
5 from loguru import logger
6 from sqlalchemy import JSON
7 from sqlalchemy import Column
8 from sqlalchemy import Float
9 from sqlalchemy import LargeBinary
10 from sqlalchemy import Text
11 from sqlalchemy import UniqueConstraint
12 from sqlalchemy.dialects.mysql import LONGTEXT
13 from sqlmodel import Field
14 from sqlmodel import Relationship
15 from sqlmodel import SQLModel
16
17
18 class Customers(SQLModel, table=True):
19 id: Optional[int] = Field(primary_key=True)
20 customer_code: str = Field(sa_column_kwargs={"index": True}, max_length=50, nullable=False)
21 parent_customer_code: Optional[str] = Field(max_length=11)
22 customer_name: str = Field(max_length=50, nullable=False)
23 contact_last_name: Optional[str] = Field(max_length=50)
24 contact_first_name: Optional[str] = Field(max_length=50)
25 phone: Optional[str] = Field(max_length=50)
26 address_line1: Optional[str] = Field(max_length=1024)
27 address_line2: Optional[str] = Field(max_length=1024)
28 city: Optional[str] = Field(max_length=50)
29 state: Optional[str] = Field(max_length=50)
30 postal_code: Optional[str] = Field(max_length=15)
31 country: Optional[str] = Field(max_length=50)
32 customer_type: Optional[str] = Field(max_length=50)
33 logo_file: Optional[str] = Field(max_length=64)
34 created_at: datetime = Field(default=datetime.utcnow())
35
36 agents: list["Agents"] = Relationship(back_populates="customer")
37 meta: Optional["CustomersMeta"] = Relationship(back_populates="customer")
38
39 def update_from_model(self, customer):
40 self.customer_code = customer.customer_code
41 self.parent_customer_code = customer.parent_customer_code
42 self.customer_name = customer.customer_name
43 self.contact_last_name = customer.contact_last_name
44 self.contact_first_name = customer.contact_first_name
45 self.phone = customer.phone
46 self.address_line1 = customer.address_line1
47 self.address_line2 = customer.address_line2
48 self.city = customer.city
49 self.state = customer.state
50 self.postal_code = customer.postal_code
51 self.country = customer.country
52 self.customer_type = customer.customer_type
53 self.logo_file = customer.logo_file
54
55
56 class CustomersMeta(SQLModel, table=True):
57 id: Optional[int] = Field(primary_key=True)
58 customer_code: str = Field(foreign_key="customers.customer_code", nullable=False)
59 customer_name: str = Field(max_length=255)
60 customer_meta_graylog_index: str = Field(max_length=1024)
61 customer_meta_graylog_stream: str = Field(max_length=1024)
62 customer_meta_grafana_org_id: str = Field(max_length=1024)
63 customer_meta_wazuh_group: str = Field(max_length=1024)
64 customer_meta_index_retention: Optional[str] = Field()
65 customer_meta_wazuh_registration_port: Optional[str] = Field()
66 customer_meta_wazuh_log_ingestion_port: Optional[str] = Field()
67 customer_meta_wazuh_api_port: Optional[str] = Field()
68 customer_meta_wazuh_auth_password: Optional[str] = Field(max_length=1024)
69 customer_meta_iris_customer_id: Optional[int] = Field()
70 customer_meta_office365_organization_id: Optional[str] = Field(max_length=1024)
71 customer_meta_portainer_stack_id: Optional[int] = Field()
72
73 # Link back to Customers
74 customer: Optional["Customers"] = Relationship(back_populates="meta")
75
76 def update_from_model(self, customer_meta):
77 if hasattr(customer_meta, "customer_code"):
78 self.customer_code = customer_meta.customer_code
79 if hasattr(customer_meta, "customer_name"):
80 self.customer_name = customer_meta.customer_name
81 self.customer_meta_graylog_index = customer_meta.customer_meta_graylog_index
82 self.customer_meta_graylog_stream = customer_meta.customer_meta_graylog_stream
83 self.customer_meta_grafana_org_id = customer_meta.customer_meta_grafana_org_id
84 self.customer_meta_wazuh_group = customer_meta.customer_meta_wazuh_group
85 self.customer_meta_index_retention = customer_meta.customer_meta_index_retention
86 self.customer_meta_wazuh_registration_port = customer_meta.customer_meta_wazuh_registration_port
87 self.customer_meta_wazuh_log_ingestion_port = customer_meta.customer_meta_wazuh_log_ingestion_port
88 self.customer_meta_wazuh_api_port = customer_meta.customer_meta_wazuh_api_port
89 self.customer_meta_wazuh_auth_password = customer_meta.customer_meta_wazuh_auth_password
90 self.customer_meta_iris_customer_id = customer_meta.customer_meta_iris_customer_id
91 self.customer_meta_office365_organization_id = customer_meta.customer_meta_office365_organization_id
92 self.customer_meta_portainer_stack_id = customer_meta.customer_meta_portainer_stack_id
93
94
95 class Agents(SQLModel, table=True):
96 id: Optional[int] = Field(primary_key=True)
97 agent_id: str = Field(index=True, max_length=256)
98 ip_address: str = Field(max_length=256)
99 os: str = Field(max_length=256)
100 hostname: str = Field(max_length=256)
101 label: str = Field(max_length=256)
102 critical_asset: bool = Field(default=False)
103 wazuh_last_seen: datetime
104 velociraptor_id: Optional[str] = Field(max_length=256)
105 velociraptor_last_seen: Optional[datetime]
106 wazuh_agent_version: str = Field(max_length=256)
107 wazuh_agent_status: str = Field("not found", max_length=256)
108 velociraptor_agent_version: Optional[str] = Field(max_length=256)
109 customer_code: Optional[str] = Field(foreign_key="customers.customer_code", max_length=256)
110 quarantined: bool = Field(default=False)
111 velociraptor_org: Optional[str] = Field(max_length=256)
112
113 customer: Optional[Customers] = Relationship(back_populates="agents")
114 vulnerabilities: Optional[list["AgentVulnerabilities"]] = Relationship(back_populates="agent")
115 data_store: list["AgentDataStore"] = Relationship(back_populates="agent")
116
117 @classmethod
118 def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code):
119 # Check if agent_last_seen is 'Unknown' and set wazuh_last_seen accordingly
120 if wazuh_agent.agent_last_seen in ("Unknown", "1970-01-01T00:00:00+00:00"):
121 wazuh_last_seen_value = datetime.strptime(
122 "1970-01-01T00:00:00+00:00",
123 "%Y-%m-%dT%H:%M:%S%z",
124 ).replace(tzinfo=None)
125 else:
126 wazuh_last_seen_value = wazuh_agent.agent_last_seen_as_datetime
127
128 return cls(
129 agent_id=wazuh_agent.agent_id,
130 hostname=wazuh_agent.agent_name,
131 ip_address=wazuh_agent.agent_ip,
132 os=wazuh_agent.agent_os,
133 label=wazuh_agent.agent_label,
134 wazuh_last_seen=wazuh_last_seen_value,
135 wazuh_agent_version=wazuh_agent.wazuh_agent_version,
136 wazuh_agent_status=wazuh_agent.wazuh_agent_status if wazuh_agent.wazuh_agent_status else "not found",
137 velociraptor_id=velociraptor_agent.client_id if velociraptor_agent and velociraptor_agent.client_id else None,
138 velociraptor_last_seen=velociraptor_agent.client_last_seen_as_datetime
139 if velociraptor_agent and velociraptor_agent.client_last_seen_as_datetime
140 else None,
141 velociraptor_agent_version=velociraptor_agent.client_version
142 if velociraptor_agent and velociraptor_agent.client_version
143 else None,
144 customer_code=customer_code,
145 velociraptor_org=velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None,
146 )
147
148 @classmethod
149 def create_wazuh_agent_from_model(cls, wazuh_agent, customer_code):
150 if wazuh_agent.agent_last_seen in ("Unknown", "1970-01-01T00:00:00+00:00"):
151 wazuh_last_seen_value = datetime.strptime(
152 "1970-01-01T00:00:00+00:00",
153 "%Y-%m-%dT%H:%M:%S%z",
154 ).replace(tzinfo=None)
155 else:
156 wazuh_last_seen_value = wazuh_agent.agent_last_seen_as_datetime
157
158 return cls(
159 agent_id=wazuh_agent.agent_id,
160 hostname=wazuh_agent.agent_name,
161 ip_address=wazuh_agent.agent_ip,
162 os=wazuh_agent.agent_os,
163 label=wazuh_agent.agent_label,
164 wazuh_last_seen=wazuh_last_seen_value,
165 wazuh_agent_version=wazuh_agent.wazuh_agent_version,
166 wazuh_agent_status=wazuh_agent.wazuh_agent_status if wazuh_agent.wazuh_agent_status else "not found",
167 customer_code=customer_code,
168 )
169
170 def update_from_model(self, wazuh_agent, velociraptor_agent, customer_code):
171 if wazuh_agent.agent_last_seen == "Unknown" or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00":
172 wazuh_last_seen_value = datetime.strptime(
173 "1970-01-01T00:00:00+00:00",
174 "%Y-%m-%dT%H:%M:%S%z",
175 ) # default datetime value
176 else:
177 wazuh_last_seen_value = wazuh_agent.agent_last_seen_as_datetime
178
179 self.agent_id = wazuh_agent.agent_id
180 self.hostname = wazuh_agent.agent_name
181 self.ip_address = wazuh_agent.agent_ip
182 self.os = wazuh_agent.agent_os
183 self.label = wazuh_agent.agent_label
184 self.wazuh_last_seen = wazuh_last_seen_value
185 self.wazuh_agent_version = wazuh_agent.wazuh_agent_version
186 self.wazuh_agent_status = wazuh_agent.wazuh_agent_status if wazuh_agent.wazuh_agent_status else "not found"
187 self.velociraptor_id = velociraptor_agent.client_id if velociraptor_agent and velociraptor_agent.client_id else None
188 self.velociraptor_last_seen = (
189 velociraptor_agent.client_last_seen_as_datetime
190 if velociraptor_agent and velociraptor_agent.client_last_seen_as_datetime
191 else None
192 )
193 self.velociraptor_agent_version = (
194 velociraptor_agent.client_version if velociraptor_agent and velociraptor_agent.client_version else None
195 )
196 self.customer_code = customer_code
197 self.velociraptor_org = velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None
198
199 def update_wazuh_agent_from_model(self, wazuh_agent, customer_code):
200 if wazuh_agent.agent_last_seen == "Unknown" or wazuh_agent.agent_last_seen == "1970-01-01T00:00:00+00:00":
201 wazuh_last_seen_value = datetime.strptime(
202 "1970-01-01T00:00:00+00:00",
203 "%Y-%m-%dT%H:%M:%S%z",
204 )
205 else:
206 wazuh_last_seen_value = wazuh_agent.agent_last_seen_as_datetime
207
208 self.agent_id = wazuh_agent.agent_id
209 self.hostname = wazuh_agent.agent_name
210 self.ip_address = wazuh_agent.agent_ip
211 self.os = wazuh_agent.agent_os
212 self.label = wazuh_agent.agent_label
213 self.wazuh_last_seen = wazuh_last_seen_value
214 self.wazuh_agent_version = wazuh_agent.wazuh_agent_version
215 self.wazuh_agent_status = wazuh_agent.wazuh_agent_status if wazuh_agent.wazuh_agent_status else "not found"
216 self.customer_code = customer_code
217
218 def update_velociraptor_details(self, velociraptor_agent):
219 logger.info(f"Updating Velociraptor details for agent {self}")
220 self.velociraptor_id = velociraptor_agent.client_id if velociraptor_agent and velociraptor_agent.client_id else None
221 self.velociraptor_last_seen = (
222 velociraptor_agent.client_last_seen_as_datetime
223 if velociraptor_agent and velociraptor_agent.client_last_seen_as_datetime
224 else None
225 )
226 self.velociraptor_agent_version = (
227 velociraptor_agent.client_version if velociraptor_agent and velociraptor_agent.client_version else None
228 )
229 logger.info(f"Updated with Velociraptor details: {self}")
230 self.velociraptor_org = velociraptor_agent.client_org if velociraptor_agent and velociraptor_agent.client_org else None
231
232
233 class AgentDataStore(SQLModel, table=True):
234 __tablename__ = "agent_datastore"
235
236 id: Optional[int] = Field(primary_key=True)
237
238 # Agent information
239 agent_id: str = Field(foreign_key="agents.agent_id", max_length=256, index=True, nullable=False)
240 velociraptor_id: str = Field(max_length=256, nullable=False)
241 # Removed customer_code - access via agent.customer_code relationship
242
243 # Artifact collection details
244 artifact_name: str = Field(max_length=255, nullable=False, index=True)
245 flow_id: str = Field(max_length=255, nullable=False, index=True)
246 collection_time: datetime = Field(default_factory=datetime.utcnow, index=True)
247
248 # MinIO storage details
249 bucket_name: str = Field(max_length=255, nullable=False, default="velociraptor-artifacts")
250 object_key: str = Field(max_length=1024, nullable=False) # Path: agent_id/flow_id/filename.zip
251 file_name: str = Field(max_length=255, nullable=False) # Original file name
252 content_type: str = Field(max_length=100, default="application/zip")
253 file_size: int = Field(nullable=False) # File size in bytes
254 file_hash: str = Field(max_length=128, nullable=False) # SHA-256 hash
255
256 # Metadata
257 uploaded_by: Optional[int] = Field(default=None) # User ID who initiated the collection
258 notes: Optional[str] = Field(sa_column=Column(Text, nullable=True))
259
260 # Status tracking
261 status: str = Field(max_length=50, default="completed", index=True) # completed, failed, processing
262 error_message: Optional[str] = Field(sa_column=Column(Text, nullable=True))
263
264 # Relationship to Agents table
265 agent: Optional["Agents"] = Relationship(back_populates="data_store")
266
267
268 class LogEntry(SQLModel, table=True):
269 __tablename__ = "log_entries"
270 id: Optional[int] = Field(primary_key=True)
271 timestamp: datetime = Field(default_factory=datetime.utcnow)
272 event_type: str = Field(default="Info", max_length=256)
273 user_id: int = Field(default=None, nullable=True)
274 route: str = Field(default=None, nullable=True, max_length=256)
275 method: str = Field(default=None, nullable=True, max_length=256)
276 status_code: int
277 message: str = Field(default=None, nullable=True, max_length=5024)
278 additional_info: str = Field(default=None, nullable=True, max_length=5024)
279
280
281 class License(SQLModel, table=True):
282 __tablename__ = "license"
283 id: Optional[int] = Field(primary_key=True)
284 license_key: str = Field(max_length=1024)
285 customer_name: str = Field(max_length=1024)
286 customer_email: str = Field(max_length=1024)
287 company_name: str = Field(max_length=1024)
288
289
290 class LicenseCache(SQLModel, table=True):
291 __tablename__ = "license_cache"
292 id: Optional[int] = Field(primary_key=True)
293 license_key: str = Field(max_length=1024, index=True)
294 feature_name: str = Field(max_length=256, index=True)
295 is_enabled: bool = Field(default=False)
296 cached_at: datetime = Field(default=datetime.utcnow, index=True)
297 expires_at: datetime = Field(index=True)
298 license_data: Optional[str] = Field(max_length=5000) # Store full license JSON as string for reference
299
300
301 class SchedulerJob(SQLModel, table=True):
302 id: str = Field(default=None, primary_key=True, nullable=False, max_length=255)
303 next_run_time: float = Field(sa_column=Column(Float(), index=True))
304 job_state: bytes = Field(sa_column=Column(LargeBinary(), nullable=False))
305
306 def __repr__(self):
307 return f"<SchedulerJob(id={self.id}, next_run_time={self.next_run_time})>"
308
309
310 class AgentVulnerabilities(SQLModel, table=True):
311 __tablename__ = "agent_vulnerabilities"
312
313 id: Optional[int] = Field(primary_key=True)
314 cve_id: str = Field(default="UNKNOWN_CVE", max_length=50, index=True)
315 severity: str = Field(default="UNKNOWN", max_length=50, index=True)
316 title: str = Field(max_length=255)
317 references: str = Field(default=None, max_length=2048)
318 status: str = Field(default="Active", max_length=50, index=True)
319 discovered_at: datetime = Field(index=True)
320 remediated_at: Optional[datetime] = Field(default=None)
321 epss_score: Optional[str] = Field(default=None, max_length=50)
322 epss_percentile: Optional[str] = Field(default=None, max_length=50)
323 package_name: Optional[str] = Field(default=None, max_length=255)
324
325 # Foreign keys
326 agent_id: str = Field(foreign_key="agents.agent_id", max_length=256, index=True)
327 customer_code: Optional[str] = Field(foreign_key="customers.customer_code", max_length=50, index=True)
328
329 # Relationship back to the Agents model
330 agent: Optional["Agents"] = Relationship(back_populates="vulnerabilities")
331
332 def update_from_model(self, vulnerability_data):
333 """Update vulnerability from external data source"""
334 if hasattr(vulnerability_data, "cve_id"):
335 self.cve_id = vulnerability_data.cve_id
336 if hasattr(vulnerability_data, "severity"):
337 self.severity = vulnerability_data.severity
338 if hasattr(vulnerability_data, "title"):
339 self.title = vulnerability_data.title
340 if hasattr(vulnerability_data, "references"):
341 self.references = vulnerability_data.references
342 if hasattr(vulnerability_data, "detected_at"):
343 self.discovered_at = vulnerability_data.detected_at
344 if hasattr(vulnerability_data, "status"):
345 self.status = vulnerability_data.status
346 if hasattr(vulnerability_data, "epss_score"):
347 self.epss_score = vulnerability_data.epss_score
348 if hasattr(vulnerability_data, "epss_percentile"):
349 self.epss_percentile = vulnerability_data.epss_percentile
350 if hasattr(vulnerability_data, "package_name"):
351 self.package_name = vulnerability_data.package_name
352 if hasattr(vulnerability_data, "remediated_at"):
353 self.remediated_at = vulnerability_data.remediated_at
354
355 @classmethod
356 def create_from_model(cls, vulnerability_data, agent_id, customer_code=None):
357 """Create a new vulnerability record from external data"""
358 return cls(
359 cve_id=getattr(vulnerability_data, "cve_id", "UNKNOWN_CVE"),
360 severity=getattr(vulnerability_data, "severity", "UNKNOWN"),
361 title=getattr(vulnerability_data, "title", ""),
362 references=getattr(vulnerability_data, "references", None),
363 status=getattr(vulnerability_data, "status", "Active"),
364 epss_score=getattr(vulnerability_data, "epss_score", None),
365 epss_percentile=getattr(vulnerability_data, "epss_percentile", None),
366 package_name=getattr(vulnerability_data, "package_name", None),
367 discovered_at=getattr(vulnerability_data, "detected_at", datetime.utcnow()),
368 agent_id=agent_id,
369 customer_code=customer_code,
370 )
371
372
373 class CustomerPortalSettings(SQLModel, table=True):
374 __tablename__ = "customer_portal_settings"
375
376 id: Optional[int] = Field(primary_key=True)
377 title: str = Field(max_length=255, default="CoPilot")
378 logo_base64: Optional[str] = Field(default=None, sa_column=Column(LONGTEXT)) # Use TEXT column for large base64 data
379 logo_mime_type: Optional[str] = Field(default=None, max_length=50) # e.g., "image/png", "image/jpeg"
380 updated_at: datetime = Field(default_factory=datetime.utcnow)
381 updated_by: Optional[int] = Field(default=None) # User ID who last updated
382
383 def update_from_request(
384 self,
385 title: Optional[str] = None,
386 logo_base64: Optional[str] = None,
387 logo_mime_type: Optional[str] = None,
388 user_id: Optional[int] = None,
389 ) -> None:
390 """
391 Update settings from request data.
392 If a field is explicitly None, restore it to default value.
393 """
394 # Get default values for restoration
395 defaults = self.get_default_values()
396
397 # Update title - if None is passed, restore to default
398 if title is not None:
399 self.title = title
400 elif title is None and hasattr(self, "_explicit_none_title"):
401 self.title = defaults["title"]
402
403 # Update logo_base64 - if None is passed, restore to default
404 if logo_base64 is not None:
405 self.logo_base64 = logo_base64
406 elif logo_base64 is None and hasattr(self, "_explicit_none_logo"):
407 self.logo_base64 = defaults["logo_base64"]
408
409 # Update logo_mime_type - if None is passed, restore to default
410 if logo_mime_type is not None:
411 self.logo_mime_type = logo_mime_type
412 elif logo_mime_type is None and hasattr(self, "_explicit_none_mime"):
413 self.logo_mime_type = defaults["logo_mime_type"]
414
415 self.updated_by = user_id
416 self.updated_at = datetime.now()
417
418 @staticmethod
419 def get_default_values() -> dict:
420 """Get default values for restoration."""
421 return {
422 "title": "CoPilot",
423 "logo_base64": None,
424 "logo_mime_type": None,
425 }
426
427 @classmethod
428 def create_default(cls) -> "CustomerPortalSettings":
429 """Create default settings."""
430 defaults = cls.get_default_values()
431 return cls(
432 title=defaults["title"],
433 logo_base64=defaults["logo_base64"],
434 logo_mime_type=defaults["logo_mime_type"],
435 )
436
437
438 class VulnerabilityReport(SQLModel, table=True):
439 __tablename__ = "vulnerability_reports"
440
441 id: Optional[int] = Field(primary_key=True)
442
443 # Report metadata
444 report_name: str = Field(max_length=255, nullable=False)
445 customer_code: str = Field(foreign_key="customers.customer_code", max_length=50, index=True, nullable=False)
446
447 # MinIO storage details
448 bucket_name: str = Field(max_length=255, nullable=False, default="vulnerability-reports")
449 object_key: str = Field(max_length=1024, nullable=False) # Path: customer_code/report_name_timestamp.csv
450 file_name: str = Field(max_length=255, nullable=False) # CSV filename
451 file_size: int = Field(nullable=False) # File size in bytes
452 file_hash: str = Field(max_length=128, nullable=False) # SHA-256 hash
453
454 # Report generation details
455 generated_at: datetime = Field(default_factory=datetime.utcnow, index=True)
456 generated_by: int = Field(nullable=False) # User ID who generated the report
457
458 # Report filters applied
459 filters_json: Optional[str] = Field(sa_column=Column(Text, nullable=True)) # JSON string of filters used
460
461 # Statistics
462 total_vulnerabilities: int = Field(default=0)
463 critical_count: int = Field(default=0)
464 high_count: int = Field(default=0)
465 medium_count: int = Field(default=0)
466 low_count: int = Field(default=0)
467
468 # Status
469 status: str = Field(max_length=50, default="completed", index=True) # completed, failed, processing
470 error_message: Optional[str] = Field(sa_column=Column(Text, nullable=True))
471
472 # Relationship to Customers table
473 customer: Optional["Customers"] = Relationship()
474
475
476 class SCAReport(SQLModel, table=True):
477 __tablename__ = "sca_reports"
478
479 id: Optional[int] = Field(primary_key=True)
480
481 # Report metadata
482 report_name: str = Field(max_length=255, nullable=False)
483 customer_code: str = Field(foreign_key="customers.customer_code", max_length=50, index=True, nullable=False)
484
485 # MinIO storage details
486 bucket_name: str = Field(max_length=255, nullable=False, default="sca-reports")
487 object_key: str = Field(max_length=1024, nullable=False) # Path: customer_code/report_name_timestamp.csv
488 file_name: str = Field(max_length=255, nullable=False) # CSV filename
489 file_size: int = Field(nullable=False, default=0) # File size in bytes
490 file_hash: str = Field(max_length=128, nullable=False, default="pending") # SHA-256 hash
491
492 # Report generation details
493 generated_at: datetime = Field(default_factory=datetime.utcnow, index=True)
494 generated_by: int = Field(nullable=False) # User ID who generated the report
495
496 # Report filters applied
497 filters_json: Optional[str] = Field(sa_column=Column(Text, nullable=True)) # JSON string of filters used
498
499 # SCA Statistics
500 total_policies: int = Field(default=0) # Number of policy results in report
501 total_checks: int = Field(default=0) # Sum of all checks across policies
502 passed_count: int = Field(default=0) # Sum of passed checks
503 failed_count: int = Field(default=0) # Sum of failed checks
504 invalid_count: int = Field(default=0) # Sum of invalid/not applicable checks
505
506 # Status tracking (for background generation)
507 status: str = Field(max_length=50, default="processing", index=True) # processing, completed, failed
508 error_message: Optional[str] = Field(sa_column=Column(Text, nullable=True))
509
510 # Relationship to Customers table
511 customer: Optional["Customers"] = Relationship()
512
513
514 class EventSources(SQLModel, table=True):
515 __tablename__ = "event_sources"
516
517 id: Optional[int] = Field(primary_key=True)
518 customer_code: str = Field(
519 foreign_key="customers.customer_code",
520 max_length=50,
521 index=True,
522 nullable=False,
523 )
524 name: str = Field(max_length=255, nullable=False)
525 index_pattern: str = Field(max_length=1024, nullable=False)
526 event_type: str = Field(max_length=50, nullable=False) # EDR, EPP, Cloud Integration, Network Security
527 time_field: str = Field(max_length=255, nullable=False, default="timestamp")
528 enabled: bool = Field(default=True)
529 # List of {key, label, width?} dicts. NULL/empty means "use the frontend's
530 # hardcoded defaults" so behaviour for un-customised sources is unchanged.
531 displayed_columns: Optional[List[dict]] = Field(
532 default=None,
533 sa_column=Column(JSON, nullable=True),
534 )
535 created_at: datetime = Field(default_factory=datetime.utcnow)
536 updated_at: datetime = Field(default_factory=datetime.utcnow)
537
538 customer: Optional["Customers"] = Relationship()
539
540 def update_from_model(self, source_data):
541 """Apply only the fields the caller explicitly set on source_data.
542
543 Pydantic 2's `model_dump(exclude_unset=True)` filters to fields the
544 client actually sent, so partial updates (e.g. PUT with just
545 `displayed_columns`) don't clobber unrelated columns to NULL. The
546 same dump recursively flattens nested Pydantic models (DisplayColumn)
547 into plain dicts, which is what the JSON column requires —
548 SQLAlchemy's json_serializer otherwise raises "Object of type
549 DisplayColumn is not JSON serializable" on commit.
550 """
551 data = source_data.model_dump(exclude_unset=True) if hasattr(source_data, "model_dump") else {}
552 for field in ("name", "index_pattern", "event_type", "time_field", "enabled", "displayed_columns"):
553 if field in data:
554 setattr(self, field, data[field])
555 self.updated_at = datetime.utcnow()
556
557
558 class EnabledDashboards(SQLModel, table=True):
559 __tablename__ = "enabled_dashboards"
560 __table_args__ = (
561 UniqueConstraint(
562 "customer_code",
563 "event_source_id",
564 "library_card",
565 "template_id",
566 name="uq_enabled_dashboard",
567 ),
568 )
569
570 id: Optional[int] = Field(primary_key=True)
571 customer_code: str = Field(
572 foreign_key="customers.customer_code",
573 max_length=50,
574 index=True,
575 nullable=False,
576 )
577 event_source_id: int = Field(foreign_key="event_sources.id", nullable=False, index=True)
578 library_card: str = Field(max_length=255, nullable=False)
579 template_id: str = Field(max_length=255, nullable=False)
580 display_name: str = Field(max_length=255, nullable=False)
581 created_at: datetime = Field(default_factory=datetime.utcnow)
582
583 # Relationships
584 customer: Optional["Customers"] = Relationship()
585 event_source: Optional["EventSources"] = Relationship()
586
587
588 class AiAnalystJob(SQLModel, table=True):
589 __tablename__ = "ai_analyst_job"
590
591 id: str = Field(primary_key=True, max_length=64)
592 alert_id: int = Field(nullable=False, index=True)
593 customer_code: str = Field(foreign_key="customers.customer_code", max_length=64, index=True, nullable=False)
594 status: str = Field(default="pending", max_length=50, index=True) # pending, running, completed, failed
595 alert_type: Optional[str] = Field(default=None, max_length=64)
596 triggered_by: str = Field(max_length=50, nullable=False) # scheduled, manual, webhook
597 template_used: Optional[str] = Field(default=None, max_length=128)
598 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
599 started_at: Optional[datetime] = Field(default=None)
600 completed_at: Optional[datetime] = Field(default=None)
601 error_message: Optional[str] = Field(sa_column=Column(Text), default=None)
602
603 customer: Optional["Customers"] = Relationship()
604 reports: list["AiAnalystReport"] = Relationship(back_populates="job")
605
606
607 class AiAnalystReport(SQLModel, table=True):
608 __tablename__ = "ai_analyst_report"
609
610 id: Optional[int] = Field(primary_key=True)
611 job_id: str = Field(foreign_key="ai_analyst_job.id", max_length=64, nullable=False, index=True)
612 alert_id: int = Field(nullable=False, index=True)
613 customer_code: str = Field(foreign_key="customers.customer_code", max_length=64, index=True, nullable=False)
614 severity_assessment: Optional[str] = Field(default=None, max_length=50) # Critical, High, Medium, Low, Informational
615 report_markdown: Optional[str] = Field(sa_column=Column(LONGTEXT), default=None)
616 summary: Optional[str] = Field(sa_column=Column(Text), default=None)
617 recommended_actions: Optional[str] = Field(sa_column=Column(Text), default=None)
618 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
619
620 job: Optional["AiAnalystJob"] = Relationship(back_populates="reports")
621 iocs: list["AiAnalystIoc"] = Relationship(back_populates="report")
622 customer: Optional["Customers"] = Relationship()
623
624
625 class AiAnalystIoc(SQLModel, table=True):
626 __tablename__ = "ai_analyst_ioc"
627
628 id: Optional[int] = Field(primary_key=True)
629 report_id: int = Field(foreign_key="ai_analyst_report.id", nullable=False, index=True)
630 alert_id: int = Field(nullable=False, index=True)
631 customer_code: str = Field(foreign_key="customers.customer_code", max_length=64, index=True, nullable=False)
632 ioc_value: str = Field(max_length=512, nullable=False)
633 ioc_type: str = Field(max_length=50, nullable=False) # ip, domain, hash, process, url, user, command
634 vt_verdict: str = Field(default="unknown", max_length=50) # malicious, suspicious, clean, unknown
635 vt_score: Optional[str] = Field(default=None, max_length=32)
636 details: Optional[str] = Field(sa_column=Column(Text), default=None)
637 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
638
639 report: Optional["AiAnalystReport"] = Relationship(back_populates="iocs")
640 customer: Optional["Customers"] = Relationship()
641 ioc_reviews: list["AiAnalystIocReview"] = Relationship(back_populates="ioc")
642
643
644 class AiAnalystReview(SQLModel, table=True):
645 __tablename__ = "ai_analyst_review"
646 __table_args__ = (
647 UniqueConstraint(
648 "report_id",
649 "reviewer_user_id",
650 name="uq_ai_analyst_review_report_reviewer",
651 ),
652 )
653
654 id: Optional[int] = Field(primary_key=True)
655 report_id: int = Field(foreign_key="ai_analyst_report.id", nullable=False, index=True)
656 alert_id: int = Field(nullable=False, index=True)
657 customer_code: str = Field(foreign_key="customers.customer_code", max_length=64, index=True, nullable=False)
658 reviewer_user_id: int = Field(nullable=False, index=True)
659 overall_verdict: Optional[str] = Field(default=None, max_length=4) # up, down
660 template_choice: Optional[str] = Field(default=None, max_length=7) # correct, wrong, partial
661 template_used: Optional[str] = Field(default=None, max_length=128)
662 rating_instructions: Optional[int] = Field(default=None) # 1–5
663 rating_artifacts: Optional[int] = Field(default=None) # 1–5
664 rating_severity: Optional[int] = Field(default=None) # 1–5
665 missing_steps: Optional[str] = Field(sa_column=Column(Text), default=None)
666 suggested_edits: Optional[str] = Field(sa_column=Column(Text), default=None)
667 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
668 updated_at: Optional[datetime] = Field(default=None)
669
670 report: Optional["AiAnalystReport"] = Relationship()
671 customer: Optional["Customers"] = Relationship()
672 ioc_reviews: list["AiAnalystIocReview"] = Relationship(back_populates="review")
673 palace_lessons: list["AiAnalystPalaceLesson"] = Relationship(back_populates="review")
674
675
676 class AiAnalystIocReview(SQLModel, table=True):
677 __tablename__ = "ai_analyst_ioc_review"
678
679 id: Optional[int] = Field(primary_key=True)
680 review_id: int = Field(foreign_key="ai_analyst_review.id", nullable=False, index=True)
681 ioc_id: int = Field(foreign_key="ai_analyst_ioc.id", nullable=False, index=True)
682 verdict_correct: bool = Field(nullable=False)
683 note: Optional[str] = Field(sa_column=Column(Text), default=None)
684 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
685
686 review: Optional["AiAnalystReview"] = Relationship(back_populates="ioc_reviews")
687 ioc: Optional["AiAnalystIoc"] = Relationship(back_populates="ioc_reviews")
688
689
690 class AiAnalystPalaceLesson(SQLModel, table=True):
691 __tablename__ = "ai_analyst_palace_lesson"
692
693 id: Optional[int] = Field(primary_key=True)
694 review_id: Optional[int] = Field(foreign_key="ai_analyst_review.id", default=None, index=True) # nullable — can be standalone
695 customer_code: str = Field(foreign_key="customers.customer_code", max_length=64, index=True, nullable=False)
696 lesson_type: str = Field(max_length=20, nullable=False) # environment, false_positives, assets, threat_intel
697 lesson_text: str = Field(sa_column=Column(Text, nullable=False))
698 durability: str = Field(default="durable", max_length=8) # one_off, durable
699 status: str = Field(default="pending", max_length=8, index=True) # pending, ingested, failed, expired
700 # drawer_id returned by mempalace add_drawer — required to call
701 # delete_drawer later when the durability sweeper expires one-offs.
702 # Nullable because legacy rows predate this column and because the
703 # drainer may fail to capture it if NanoClaw returns a malformed body.
704 drawer_id: Optional[str] = Field(default=None, max_length=64, index=True)
705 ingested_at: Optional[datetime] = Field(default=None)
706 # Timestamp of the sweeper's delete_drawer call. Set when status flips
707 # from 'ingested' → 'expired' so audit queries can tell "never swept"
708 # apart from "swept but failed".
709 expired_at: Optional[datetime] = Field(default=None)
710 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
711
712 review: Optional["AiAnalystReview"] = Relationship(back_populates="palace_lessons")
713 customer: Optional["Customers"] = Relationship()
714
715
716 # ---------------------------------------------------------------------------
717 # Notification routing
718 #
719 # Per-customer "where do we tell someone about an investigation result"
720 # config. Phase 1 ships with two delivery channels — Slack incoming
721 # webhooks and SMTP email — and is intentionally provider-direct (no
722 # Shuffle dependency yet). Phase 2 adds a `customer_shuffle_integrations`
723 # table and an `integration_id` FK on the routes table to layer Shuffle's
724 # 3,000+ app catalog on top, without breaking the Phase 1 routes.
725 #
726 # Triggers and severities are stored as plain strings (not enums) on
727 # purpose — adding a new trigger or severity tier later is a data-only
728 # change, no migration. The CRUD layer enforces the v1 set.
729 # ---------------------------------------------------------------------------
730
731
732 class CustomerNotificationRoute(SQLModel, table=True):
733 __tablename__ = "customer_notification_route"
734
735 id: Optional[int] = Field(primary_key=True)
736 customer_code: str = Field(
737 foreign_key="customers.customer_code",
738 max_length=64,
739 index=True,
740 nullable=False,
741 )
742
743 # Human label shown in the UI list. Without this, users would have to
744 # mentally parse channel+destination columns to identify a rule.
745 name: str = Field(max_length=128, nullable=False)
746
747 # 'investigation_complete' (any successful investigation, regardless
748 # of verdict) or 'severity_critical_or_high' (Critical/High only).
749 # Stored as string so adding new triggers later is a data-only change.
750 trigger: str = Field(max_length=64, nullable=False, index=True)
751
752 # Always 'shuffle' on the current code path — the column kept its
753 # varchar shape so we can re-introduce direct channels (raw webhook,
754 # PagerDuty REST, etc.) later without a migration. Pairs with the
755 # shuffle_integration_id FK to scope the dispatch to a specific org.
756 channel: str = Field(max_length=32, nullable=False)
757
758 # Free-form destination hint (Slack channel, email recipient,
759 # handle, etc.) — Shuffle's app agent figures out how to route it
760 # within the authenticated app at dispatch time.
761 destination: str = Field(sa_column=Column(Text, nullable=False))
762
763 # 'Critical' | 'High' | 'Medium' | 'Low' | 'Informational'. Inclusive
764 # — a 'High' route fires on Critical and High.
765 min_severity: str = Field(max_length=20, nullable=False, default="Medium")
766
767 # Optional Jinja-style override for the dispatched message body.
768 # Default templates live in the service layer; this lets a customer
769 # tune wording without a code change. Phase 4 ships the polished
770 # default set; Phase 1 ships a no-frills fallback.
771 format_template: Optional[str] = Field(sa_column=Column(Text), default=None)
772
773 enabled: bool = Field(default=True, nullable=False)
774
775 # Denormalized for the UI list so we can show "fired 2h ago" without
776 # joining the dispatch log on every render. Maintained by the
777 # dispatch service.
778 last_dispatched_at: Optional[datetime] = Field(default=None)
779 dispatch_count: int = Field(default=0, nullable=False)
780
781 # CoPilot user who created the route — audit trail for change
782 # management. Populated from the auth context in the route handler.
783 created_by: Optional[str] = Field(default=None, max_length=128)
784
785 # ----- Phase 2: Shuffle channel routing -----
786 # Populated when channel='shuffle'. NULL for legacy SMTP routes.
787 # The (integration_id, app_id, app_name) triple together describes
788 # "which Shuffle org" + "which app inside that org" + "label for the
789 # UI." `app_id` is the Shuffle app UUID we POST to
790 # /api/v1/apps/{app_id}/mcp; `app_name` is the human-readable label
791 # (e.g. "Slack") we cache so the UI doesn't have to roundtrip to
792 # Shuffle to render the route list.
793 shuffle_integration_id: Optional[int] = Field(
794 default=None,
795 foreign_key="customer_shuffle_integration.id",
796 index=True,
797 )
798 shuffle_app_id: Optional[str] = Field(default=None, max_length=64)
799 shuffle_app_name: Optional[str] = Field(default=None, max_length=128)
800
801 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
802 updated_at: Optional[datetime] = Field(default=None)
803
804 customer: Optional["Customers"] = Relationship()
805 # NB: no `back_populates` on the reverse relationships below. The
806 # dispatch service never traverses these — but with back_populates
807 # configured, SQLAlchemy fires implicit synchronous loads on the
808 # parent collections during flush() to keep the in-session graph in
809 # sync, which throws MissingGreenlet under AsyncSession. One-way
810 # foreign keys are fine here; we walk them via explicit queries
811 # (`session.get(...)`) when we need them.
812 dispatches: list["NotificationDispatchLog"] = Relationship(sa_relationship_kwargs={"overlaps": "route"})
813 shuffle_integration: Optional["CustomerShuffleIntegration"] = Relationship(sa_relationship_kwargs={"overlaps": "routes"})
814
815
816 class NotificationDispatchLog(SQLModel, table=True):
817 __tablename__ = "notification_dispatch_log"
818 __table_args__ = (
819 # Idempotency key: re-running the same investigation must not
820 # re-fire the same notification. The dispatch service does
821 # "INSERT ... ON CONFLICT DO NOTHING" against this constraint and
822 # short-circuits if a row already exists.
823 UniqueConstraint(
824 "customer_code",
825 "alert_id",
826 "route_id",
827 "trigger",
828 name="uq_notif_dispatch_idem",
829 ),
830 )
831
832 id: Optional[int] = Field(primary_key=True)
833 customer_code: str = Field(max_length=64, index=True, nullable=False)
834 alert_id: int = Field(nullable=False, index=True)
835 route_id: int = Field(
836 foreign_key="customer_notification_route.id",
837 nullable=False,
838 index=True,
839 )
840 trigger: str = Field(max_length=64, nullable=False)
841
842 dispatched_at: datetime = Field(default_factory=datetime.utcnow, index=True)
843 # 'sent' | 'failed' | 'skipped' (skipped = filter mismatch reached
844 # the log path, e.g. a route whose enabled=false flipped during a
845 # batch). Phase 4 retry semantics will add 'retrying'.
846 status: str = Field(max_length=16, nullable=False, index=True)
847 error_message: Optional[str] = Field(sa_column=Column(Text), default=None)
848 # Wall-clock latency of the underlying provider call (Slack POST or
849 # SMTP send), excluding our own DB work. Useful for spotting flaky
850 # webhooks before they become a customer complaint.
851 latency_ms: Optional[int] = Field(default=None)
852 # First 500 chars of the formatted body. Stored for debugging — when
853 # a customer says "the message looked wrong" we want to see what we
854 # actually sent without storing the entire body history.
855 payload_preview: Optional[str] = Field(sa_column=Column(Text), default=None)
856 # Phase 2: Shuffle's POST /apps/{id}/mcp returns a fire-and-record
857 # execution_id. Stored here so an admin can pivot from "this
858 # notification didn't arrive at Slack" → look up the run in
859 # shuffler.io's UI to see whether Shuffle accepted the dispatch but
860 # the downstream app rejected it. Null for non-Shuffle channels.
861 shuffle_execution_id: Optional[str] = Field(default=None, max_length=128)
862
863 # See note on CustomerNotificationRoute.dispatches — back_populates
864 # removed deliberately to keep AsyncSession flush() synchronous-IO-free.
865 # `overlaps` makes the deliberate one-way nature explicit to SQLAlchemy 2.
866 route: Optional["CustomerNotificationRoute"] = Relationship(sa_relationship_kwargs={"overlaps": "dispatches"})
867
868
869 class CustomerShuffleIntegration(SQLModel, table=True):
870 __tablename__ = "customer_shuffle_integration"
871
872 id: Optional[int] = Field(primary_key=True)
873 customer_code: str = Field(
874 foreign_key="customers.customer_code",
875 max_length=64,
876 index=True,
877 nullable=False,
878 )
879
880 # The customer's Shuffle Org-Id. SOCfortress's deployment-wide
881 # `SHUFFLE_API_KEY` (admin-scoped, lives in the connectors table)
882 # has access to every org; the per-customer differentiator is this
883 # Org-Id, sent as the `Org-Id` header on each dispatch so Shuffle
884 # routes the call to the correct org's authenticated apps. Stored
885 # opaquely as a string — Shuffle uses a UUID format today but we
886 # don't depend on that.
887 shuffle_org_id: str = Field(max_length=64, nullable=False)
888
889 # Human label, e.g. "Acme Production Shuffle". Surfaced in the
890 # CoPilot UI's integration picker; not sent to Shuffle.
891 display_name: str = Field(max_length=128, nullable=False)
892
893 enabled: bool = Field(default=True, nullable=False)
894 # Updated by the dispatch service whenever a route referencing this
895 # integration successfully fires — useful for spotting integrations
896 # that are configured but never used.
897 last_used_at: Optional[datetime] = Field(default=None)
898
899 created_by: Optional[str] = Field(default=None, max_length=128)
900 created_at: datetime = Field(default_factory=datetime.utcnow, index=True)
901 updated_at: Optional[datetime] = Field(default=None)
902
903 customer: Optional["Customers"] = Relationship()
904 # See note on CustomerNotificationRoute.dispatches — back_populates
905 # removed for AsyncSession compatibility.
906 routes: list["CustomerNotificationRoute"] = Relationship(sa_relationship_kwargs={"overlaps": "shuffle_integration"})