main
py 308 lines 10.5 KB
Raw
1 from typing import List
2 from typing import Optional
3
4 from fastapi import HTTPException
5 from pydantic import BaseModel
6 from pydantic import Field
7 from pydantic import field_validator
8
9
10 class GrafanaOrganizations(BaseModel):
11 id: int = Field(..., description="The ID of the organization.")
12 name: str = Field(..., description="The name of the organization.")
13
14
15 class GrafanaOrganizationsResponse(BaseModel):
16 message: str = Field(..., description="The message from the response.")
17 orgs: List[GrafanaOrganizations] = Field(..., description="The organizations collected from Grafana.")
18 success: bool = Field(..., description="The success of the response.")
19
20
21 class GrafanaOrganizationDashboards(BaseModel):
22 id: int
23 uid: str
24 title: str
25 uri: str
26 url: str
27 slug: str
28 type: str
29 tags: List[str]
30 isStarred: bool
31 sortMeta: int
32 folderId: Optional[int] = None
33 folderUid: Optional[str] = None
34 folderTitle: Optional[str] = None
35 folderUrl: Optional[str] = None
36
37
38 class GrafanaDashboardResponse(BaseModel):
39 message: str = Field(..., description="The message from the response.")
40 dashboards: List[GrafanaOrganizationDashboards] = Field(..., description="The dashboards collected from Grafana.")
41 success: bool = Field(..., description="The success of the response.")
42
43
44 class Annotation(BaseModel):
45 builtIn: int
46 datasource: dict # More specific model can be created for datasource
47 enable: bool
48 hide: bool
49 iconColor: str
50 name: str
51 type: str
52
53
54 class Threshold(BaseModel):
55 color: str
56 value: Optional[float] = None
57
58
59 class FieldConfigDefaults(BaseModel):
60 color: dict # More specific model can be created for color
61 custom: dict # More specific model can be created for custom options
62 mappings: List[dict] # More specific model can be created for mappings
63 thresholds: dict # Utilize Threshold model
64
65
66 class FieldConfig(BaseModel):
67 defaults: FieldConfigDefaults
68 overrides: List[dict] # More specific model can be created for overrides
69
70
71 class GridPos(BaseModel):
72 h: int
73 w: int
74 x: int
75 y: int
76
77
78 class PanelOptions(BaseModel):
79 legend: dict # More specific model can be created for legend
80 tooltip: dict # More specific model can be created for tooltip
81
82
83 class Panel(BaseModel):
84 fieldConfig: Optional[FieldConfig] = Field(None, description="The field configuration for the panel.")
85 gridPos: Optional[GridPos] = Field(None, description="The grid position for the panel.")
86 id: int
87 options: Optional[PanelOptions] = Field(None, description="The options for the panel.")
88 title: str
89 type: Optional[str] = Field(None, description="The type of the panel.")
90 collapsed: Optional[bool] = None # Optional field for row panel
91 panels: List["Panel"] = [] # Nested list for row panel
92
93
94 class Templating(BaseModel):
95 list: List[dict] # More specific model can be created for template variables
96
97
98 class DashboardDetails(BaseModel):
99 annotations: dict # Use Annotation model for values in the list
100 editable: bool
101 fiscalYearStartMonth: int
102 graphTooltip: int
103 id: int
104 links: List[dict] # More specific model can be created for links
105 liveNow: bool
106 panels: List[Panel]
107 refresh: str
108 schemaVersion: int
109 tags: List[str]
110 templating: Templating
111 time: dict # More specific model can be created for time range
112 timepicker: dict # More specific model can be created for timepicker options
113 timezone: str
114 title: str
115 uid: str
116 version: int
117 weekStart: str
118
119
120 class MetaDetails(BaseModel):
121 type: str
122 canSave: bool
123 canEdit: bool
124 canAdmin: bool
125 canStar: bool
126 canDelete: bool
127 slug: str
128 url: str
129 expires: str
130 created: str
131 updated: str
132 updatedBy: str
133 createdBy: str
134 version: int
135 hasAcl: bool
136 isFolder: bool
137 folderId: int
138 folderUid: str
139 folderTitle: str
140 folderUrl: str
141 provisioned: bool
142 provisionedExternalId: Optional[str] = None
143 annotationsPermissions: dict # More specific model can be created for permissions
144
145
146 class GrafanaDashboardDetails(BaseModel):
147 meta: MetaDetails
148 dashboard: DashboardDetails
149
150
151 class GrafanaDashboardDetailsResponse(BaseModel):
152 message: str = Field(..., description="The message from the response.")
153 dashboard_details: GrafanaDashboardDetails = Field(..., description="The dashboard details collected from Grafana.")
154 success: bool = Field(..., description="The success of the response.")
155
156
157 class GrafanaDashboardPanelsResponse(BaseModel):
158 message: str = Field(..., description="The message from the response.")
159 panels: List[Panel] = Field(..., description="The panels collected from Grafana.")
160 success: bool = Field(..., description="The success of the response.")
161
162
163 class TimeRange(BaseModel):
164 value: int
165 unit: str
166
167 @field_validator("unit")
168 @classmethod
169 def validate_unit(cls, unit):
170 valid_units = ["m", "h", "d"]
171 if unit not in valid_units:
172 raise HTTPException(status_code=400, detail=f"Invalid time range unit: {unit}. Must be one of {valid_units}")
173 return unit
174
175
176 class GrafanaGenerateIframeLinksRequest(BaseModel):
177 org_id: int = Field(..., description="The ID of the organization.")
178 dashboard_title: str = Field(..., description="The title of the dashboard.")
179 dashboard_uid: str = Field(..., description="The UID of the dashboard.")
180 panel_id: int = Field(..., description="The IDs of the panels.")
181 time_range: TimeRange = Field(..., description="Time range in minutes, hours, or days")
182 theme: Optional[str] = Field(None, description="The theme of the panel.")
183
184
185 class GrafanaLinksList(BaseModel):
186 panel_id: int
187 panel_url: str
188
189
190 class GrafanaGenerateIframeLinksResponse(BaseModel):
191 message: str = Field(..., description="The message from the response.")
192 links: List[GrafanaLinksList] = Field(..., description="The links collected from Grafana.")
193 success: bool = Field(..., description="The success of the response.")
194
195
196 class RequestPanel(BaseModel):
197 panel_id: int = Field(..., description="Panel ID")
198 org_id: int = Field(..., description="Organization ID")
199 dashboard_title: str = Field(..., description="Dashboard title")
200 dashboard_uid: str = Field(..., description="Dashboard UID")
201 panel_url: Optional[str] = Field(None, description="Panel URL")
202 panel_base64: Optional[str] = Field(None, description="Panel Base64")
203 row_id: Optional[int] = Field(None, description="Row ID")
204 panel_width: int = Field(..., description="Panel width")
205 panel_height: int = Field(..., description="Panel height")
206 theme: Optional[str] = Field(None, description="Panel theme")
207
208
209 class RequestRow(BaseModel):
210 id: int = Field(..., description="Row ID")
211 panels: List[RequestPanel] = Field(..., description="Panels in the row")
212
213
214 class GenerateReportRequest(BaseModel):
215 company_name: str = Field(..., description="Company name", examples=["SOC Fortress"])
216 timerange_text: str = Field(..., description="Time range text", examples=["Last 7 days"])
217 logo_base64: str = Field(
218 ...,
219 description="Base64 encoded logo",
220 examples=["data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABgAAAAYCAYAAADgdz34AAABjklEQVRIS+2Vv0oDQRDG"],
221 )
222 timerange: str = Field(..., description="Time range for the report")
223 # rows: List[RequestRow] = Field(..., description="Rows in the report")
224 rows: List[RequestRow] = Field(
225 ...,
226 description="Rows in the report",
227 examples=[
228 [
229 {
230 "id": 1710437961108,
231 "panels": [
232 {
233 "panel_id": 5,
234 "org_id": 1,
235 "dashboard_title": "HUNTRESS - _SUMMARY",
236 "dashboard_uid": "ab9bab2c-5d86-43e7-bac2-c1d68fc91342",
237 "panel_width": 500,
238 "panel_height": 300,
239 },
240 {
241 "panel_id": 3,
242 "org_id": 1,
243 "dashboard_title": "HUNTRESS - _SUMMARY",
244 "dashboard_uid": "ab9bab2c-5d86-43e7-bac2-c1d68fc91342",
245 "panel_width": 500,
246 "panel_height": 300,
247 },
248 ],
249 },
250 {
251 "id": 1710437961109,
252 "panels": [
253 {
254 "panel_id": 5,
255 "org_id": 1,
256 "dashboard_title": "HUNTRESS - _SUMMARY",
257 "dashboard_uid": "ab9bab2c-5d86-43e7-bac2-c1d68fc91342",
258 "panel_width": 500,
259 "panel_height": 300,
260 },
261 {
262 "panel_id": 3,
263 "org_id": 1,
264 "dashboard_title": "HUNTRESS - _SUMMARY",
265 "dashboard_uid": "ab9bab2c-5d86-43e7-bac2-c1d68fc91342",
266 "panel_width": 500,
267 "panel_height": 300,
268 },
269 ],
270 },
271 ],
272 ],
273 )
274
275
276 class GenerateReportCreation(BaseModel):
277 urls: list[str] = Field(
278 [
279 (
280 "http://ashdevcopilot01.socfortress.local:3000/d-solo/ab9bab2c-5d86-43e7-bac2-c1d68fc91342/"
281 "huntress-summary?orgId=1&from=1708725633941&to=1709330433941&panelId=5"
282 ),
283 (
284 "http://ashdevcopilot01.socfortress.local:3000/d-solo/ab9bab2c-5d86-43e7-bac2-c1d68fc91342/"
285 "huntress-summary?orgId=1&from=1708725654862&to=1709330454862&panelId=1"
286 ),
287 (
288 "http://ashdevcopilot01.socfortress.local:3000/d-solo/a1891b09-fba9-498e-807e-1ad774c8557f/"
289 "sap-users-auth?orgId=44&from=1709303384274&to=1709389784274&panelId=43"
290 ),
291 (
292 "http://ashdevcopilot01.socfortress.local:3000/d-solo/ab9bab2c-5d86-43e7-bac2-c1d68fc91342/"
293 "huntress-summary?orgId=1&from=1706799780600&to=1709391780600&panelId=10"
294 ),
295 ],
296 description="List of URLs to generate screenshots for",
297 )
298
299
300 class Base64Image(BaseModel):
301 base64_image: str
302 url: str
303
304
305 class GenerateReportResponse(BaseModel):
306 base64_result: str
307 message: str
308 success: bool