@cryptotaxi247 / CoPilot
Code Commits Issues Pulls Sessions CI/CD
main
main
py 183 lines 5.23 KB
Raw
1 from datetime import datetime
2 from typing import Any
3 from typing import Dict
4 from typing import List
5 from typing import Optional
6
7 from pydantic import BaseModel
8 from pydantic import ConfigDict
9 from pydantic import Field
10
11
12 class WazuhVulnerabilityData(BaseModel):
13 """Schema for processing raw Wazuh vulnerability data from Elasticsearch"""
14
15 cve_id: str = Field(..., alias="id")
16 severity: str
17 title: str = Field(..., alias="description")
18 references: Optional[str] = Field(None, alias="reference")
19 detected_at: datetime
20 published_at: Optional[datetime] = None
21 base_score: Optional[float] = Field(None, alias="base")
22 package_name: Optional[str] = None
23 package_version: Optional[str] = None
24 package_architecture: Optional[str] = None
25 model_config = ConfigDict(populate_by_name=True)
26
27
28 class AgentVulnerabilityOut(BaseModel):
29 """Output schema for agent vulnerabilities"""
30
31 id: int
32 cve_id: str
33 severity: str
34 title: str
35 references: Optional[str] = None
36 status: str
37 discovered_at: datetime
38 remediated_at: Optional[datetime] = None
39 epss_score: Optional[str] = None
40 epss_percentile: Optional[str] = None
41 package_name: Optional[str] = None
42 agent_id: str
43 customer_code: Optional[str] = None
44
45
46 class AgentVulnerabilitiesResponse(BaseModel):
47 """Response schema for agent vulnerabilities"""
48
49 vulnerabilities: List[AgentVulnerabilityOut]
50 success: bool
51 message: str
52 total_count: int
53
54
55 class VulnerabilitySyncRequest(BaseModel):
56 """Request schema for syncing vulnerabilities - all fields optional"""
57
58 customer_code: Optional[str] = Field(None, description="Optional customer code filter")
59 agent_name: Optional[str] = Field(None, description="Optional specific agent name")
60 force_refresh: bool = Field(False, description="Force refresh of existing vulnerabilities")
61
62
63 class VulnerabilitySyncResponse(BaseModel):
64 """Response schema for vulnerability sync operations"""
65
66 success: bool
67 message: str
68 synced_count: int
69 errors: List[str] = []
70
71
72 class VulnerabilityStatsResponse(BaseModel):
73 """Response schema for vulnerability statistics"""
74
75 total_vulnerabilities: int
76 critical_count: int
77 high_count: int
78 medium_count: int
79 low_count: int
80 by_customer: dict = {}
81 success: bool
82 message: str
83
84
85 class VulnerabilityDeleteResponse(BaseModel):
86 """Response schema for vulnerability delete operations"""
87
88 success: bool
89 message: str
90 deleted_count: int
91 errors: List[str] = []
92
93
94 class VulnerabilitySearchRequest(BaseModel):
95 """Request schema for searching vulnerabilities from Wazuh indexer"""
96
97 customer_code: Optional[str] = Field(None, description="Filter by customer code")
98 agent_name: Optional[str] = Field(None, description="Filter by agent hostname")
99 severity: Optional[str] = Field(None, description="Filter by severity (Critical, High, Medium, Low)")
100 page: int = Field(1, description="Page number for pagination", ge=1)
101 page_size: int = Field(50, description="Number of vulnerabilities per page", ge=1, le=1000)
102 cve_id: Optional[str] = Field(None, description="Filter by specific CVE ID")
103 package_name: Optional[str] = Field(None, description="Filter by package name")
104
105
106 class VulnerabilitySearchItem(BaseModel):
107 """Individual vulnerability item from search results"""
108
109 cve_id: str
110 severity: str
111 title: str
112 agent_name: str
113 customer_code: Optional[str] = None
114 references: Optional[str] = None
115 detected_at: datetime
116 published_at: Optional[datetime] = None
117 base_score: Optional[float] = None
118 package_name: Optional[str] = None
119 package_version: Optional[str] = None
120 package_architecture: Optional[str] = None
121 epss_score: Optional[str] = None
122 epss_percentile: Optional[str] = None
123
124
125 class VulnerabilitySearchResponse(BaseModel):
126 """Response schema for vulnerability search results with pagination"""
127
128 vulnerabilities: List[VulnerabilitySearchItem]
129 total_count: int
130 critical_count: int
131 high_count: int
132 medium_count: int
133 low_count: int
134 page: int
135 page_size: int
136 total_pages: int
137 has_next: bool
138 has_previous: bool
139 success: bool
140 message: str
141 filters_applied: dict = {}
142
143
144 class VulnerabilityReportGenerateRequest(BaseModel):
145 customer_code: str
146 report_name: Optional[str] = None # Auto-generate if not provided
147 agent_name: Optional[str] = None
148 severity: Optional[str] = None
149 cve_id: Optional[str] = None
150 package_name: Optional[str] = None
151 include_epss: bool = False
152
153
154 class VulnerabilityReportResponse(BaseModel):
155 id: int
156 report_name: str
157 customer_code: str
158 file_name: str
159 file_size: int
160 generated_at: datetime
161 generated_by: int
162 total_vulnerabilities: int
163 critical_count: int
164 high_count: int
165 medium_count: int
166 low_count: int
167 filters_applied: Dict[str, Any]
168 status: str
169 download_url: Optional[str] = None
170
171
172 class VulnerabilityReportListResponse(BaseModel):
173 reports: List[VulnerabilityReportResponse]
174 total_count: int
175 success: bool
176 message: str
177
178
179 class VulnerabilityReportGenerateResponse(BaseModel):
180 success: bool
181 message: str
182 report: Optional[VulnerabilityReportResponse] = None
183 error: Optional[str] = None