main
py 407 lines 11.1 KB
Raw
1 import ipaddress
2 import re
3 from abc import ABC
4 from typing import Dict
5 from typing import Optional
6 from typing import Union
7
8 import httpx
9 import regex
10 from fastapi import HTTPException
11 from loguru import logger
12 from sqlalchemy.ext.asyncio import AsyncSession
13
14 from app.integrations.utils.schema import ShufflePayload
15 from app.utils import get_customer_alert_settings
16
17
18 #################### ! DFIR IRIS ASSET VALIDATOR ! ####################
19 class AssetValidator(ABC):
20 """
21 Base class for asset validators.
22
23 Attributes:
24 os (str): The OS to be validated.
25 """
26
27 ASSET_TYPE_ID: int = 1
28
29 def __init__(self, os: str) -> None:
30 """
31 Initialize a Validator.
32
33 Args:
34 os (str): The OS to be validated.
35 """
36 self.os = os.lower()
37
38 def validate(self) -> Dict[str, Union[bool, str, int]]:
39 """
40 Validate the OS.
41
42 If the OS matches the type of this validator,
43 the method returns a dictionary indicating success, the matching message, and the asset type id.
44
45 Returns:
46 Dict[str, Union[bool, str, int]]: The validation result.
47 """
48 raise NotImplementedError
49
50
51 class WindowsAssetValidator(AssetValidator):
52 """
53 Class to check if an OS is Windows.
54 """
55
56 ASSET_TYPE_ID = 9
57
58 def validate(self) -> Dict[str, Union[bool, str, int]]:
59 if "windows" in self.os:
60 return {
61 "success": True,
62 "message": f"{self.os} is a valid Windows OS.",
63 "asset_type_id": self.ASSET_TYPE_ID,
64 }
65 else:
66 return {
67 "success": False,
68 "message": f"{self.os} is not a Windows OS.",
69 "asset_type_id": self.ASSET_TYPE_ID,
70 }
71
72
73 class LinuxAssetValidator(AssetValidator):
74 """
75 Class to check if an OS is Linux.
76 """
77
78 ASSET_TYPE_ID = 4
79
80 def validate(self) -> Dict[str, Union[bool, str, int]]:
81 if "linux" in self.os:
82 return {
83 "success": True,
84 "message": f"{self.os} is a valid Linux OS.",
85 "asset_type_id": self.ASSET_TYPE_ID,
86 }
87 else:
88 return {
89 "success": False,
90 "message": f"{self.os} is not a Linux OS.",
91 "asset_type_id": self.ASSET_TYPE_ID,
92 }
93
94
95 class FirewallAssetValidator(AssetValidator):
96 """
97 Class to check if an OS is Firewall.
98 """
99
100 ASSET_TYPE_ID = 2
101
102 def validate(self) -> Dict[str, Union[bool, str, int]]:
103 if "firewall" in self.os:
104 return {
105 "success": True,
106 "message": f"{self.os} is a valid Firewall OS.",
107 "asset_type_id": self.ASSET_TYPE_ID,
108 }
109 else:
110 return {
111 "success": False,
112 "message": f"{self.os} is not a Firewall OS.",
113 "asset_type_id": self.ASSET_TYPE_ID,
114 }
115
116
117 class UbuntuAssetValidator(AssetValidator):
118 """
119 Class to check if an OS is Ubuntu.
120 """
121
122 ASSET_TYPE_ID = 4
123
124 def validate(self) -> Dict[str, Union[bool, str, int]]:
125 if "ubuntu" in self.os:
126 return {
127 "success": True,
128 "message": f"{self.os} is a valid Ubuntu OS.",
129 "asset_type_id": self.ASSET_TYPE_ID,
130 }
131 else:
132 return {
133 "success": False,
134 "message": f"{self.os} is not an Ubuntu OS.",
135 "asset_type_id": self.ASSET_TYPE_ID,
136 }
137
138
139 class AssetTypeResolver:
140 """
141 Class to iterate over asset validators and return the successful validator's asset type id.
142 """
143
144 def __init__(self, os: str):
145 """
146 Initialize AssetTypeResolver.
147
148 Args:
149 os (str): The OS to be validated.
150 """
151 self.os = os
152 self.validators = [
153 WindowsAssetValidator,
154 LinuxAssetValidator,
155 FirewallAssetValidator,
156 UbuntuAssetValidator,
157 ]
158
159 def get_asset_type_id(self) -> int:
160 """
161 Iterate over validators and return the successful validator's asset type id.
162
163 Returns:
164 int: The asset type id.
165 """
166 for Validator in self.validators:
167 validator = Validator(self.os)
168 result = validator.validate()
169 if result["success"] is True:
170 return result["asset_type_id"]
171
172 # Return default asset type id (1) if no validators succeed
173 return 1
174
175
176 #################### ! DFIR IRIS ASSET VALIDATOR END ! ####################
177
178
179 #################### ! DFIR IRIS IOC VALIDATOR ! ##########################
180
181
182 class IoCValidator(ABC):
183 """
184 Base class for validators.
185
186 Attributes:
187 value (str): The value to be validated.
188 """
189
190 PATTERN: Optional[str] = None # type: ignore
191 IOC_TYPE: Optional[int] = None # type: ignore
192
193 def __init__(self, value: str) -> None:
194 """
195 Initialize a Validator.
196
197 Args:
198 value (str): The value to be validated.
199 """
200 self.value = value
201
202 def validate(self) -> Dict[str, Union[bool, str, int]]:
203 """
204 Validate the value.
205
206 If the value matches the pattern,
207 the method returns a dictionary indicating success, the matching message, and the IOC type.
208
209 Returns:
210 Dict[str, Union[bool, str, int]]: The validation result.
211 """
212 logger.info(f"Validating {self.value} against {self.PATTERN}.")
213 if self.PATTERN and regex.match(self.PATTERN, self.value, re.IGNORECASE):
214 return {
215 "success": True,
216 "message": f"{self.value} matches the pattern.",
217 "ioc_type": self.IOC_TYPE,
218 }
219 else:
220 return {
221 "success": False,
222 "message": f"{self.value} does not match the pattern.",
223 "ioc_type": self.IOC_TYPE,
224 }
225
226
227 class IPv4AddressValidator(IoCValidator):
228 """
229 Class to check if a string is a valid IPv4 address.
230 """
231
232 IOC_TYPE = 76
233
234 def validate(self) -> Dict[str, Union[bool, str, int]]:
235 """
236 Validate if the given value is a valid IPv4 address.
237
238 Returns:
239 dict: A dictionary containing success status, message, and the associated IoC type.
240 """
241 try:
242 # if the value is like this `162.159.133.233|443` strip the port
243 if "|" in self.value:
244 self.value = self.value.split("|")[0]
245 logger.info(f"Validating {self.value} as an IPv4 address.")
246 ipaddress.IPv4Address(self.value)
247 return {
248 "success": True,
249 "message": f"{self.value} is a valid IPv4 address.",
250 "ioc_type": self.IOC_TYPE,
251 }
252 except ValueError:
253 return {
254 "success": False,
255 "message": f"{self.value} is not a valid IPv4 address.",
256 "ioc_type": self.IOC_TYPE,
257 }
258
259
260 class HashValidator(IoCValidator):
261 """
262 Class to check if a string is a valid SHA256 hash.
263 """
264
265 PATTERN = r"^[a-fA-F\d]{64}$"
266 IOC_TYPE = 113
267
268
269 class DomainValidator(IoCValidator):
270 """
271 Class to check if a string is a valid domain name.
272 """
273
274 PATTERN = r"^([a-z0-9]+(-[a-z0-9]+)*\.)+[a-z]{2,}$"
275 IOC_TYPE = 20
276
277
278 #################### ! DFIR IRIS IOC VALIDATOR END ! ##########################
279
280
281 async def get_asset_type_id(os: str) -> int:
282 """
283 Use AssetTypeResolver to determine the asset type ID to set within DFIR-IRIS.
284
285 Parameters
286 ----------
287 os : str
288 The operating system (OS) string used to resolve the asset type ID.
289
290 Returns
291 -------
292 int
293 The ID corresponding to the asset type.
294 """
295 asset_resolver = AssetTypeResolver(os)
296 return asset_resolver.get_asset_type_id()
297
298
299 async def validate_ioc_type(ioc_value: str) -> str:
300 """
301 Validate IoC type using validators.
302
303 Parameters
304 ----------
305 ioc_value : str
306 The value to validate the IoC type.
307
308 Returns
309 -------
310 str
311 The type of the IoC. Returns None if validation fails.
312 """
313 validators = [IPv4AddressValidator, HashValidator, DomainValidator]
314 ioc_type = None
315
316 for Validator in validators:
317 validator = Validator(ioc_value)
318 result = validator.validate()
319
320 if result["success"]:
321 ioc_type = result["ioc_type"]
322 break
323
324 if ioc_type is None:
325 logger.error("Failed to validate IoC value.")
326 return ioc_type
327
328
329 async def send_to_shuffle(payload: ShufflePayload, session: AsyncSession) -> bool:
330 """
331 Sends payload to Shuffle listening Webhook asynchronously using httpx.
332 """
333 logger.info(f"Sending {payload} to Shuffle Webhook.")
334 try:
335 async with httpx.AsyncClient(verify=False) as client:
336 response = await client.post(
337 (
338 await get_customer_alert_settings(
339 customer_code=payload.customer_code,
340 session=session,
341 )
342 ).shuffle_endpoint,
343 json=payload.to_dict(),
344 )
345
346 return response.status_code == 200
347
348 except Exception as e:
349 logger.error(f"Error: {e}")
350 raise HTTPException(
351 status_code=500,
352 detail=f"Error: {e}",
353 )
354
355
356 # def send_to_wazuh(msg) -> None:
357 # # Uncomment when doing dev work
358 # # logger.info(f"Sending {msg} to Wazuh Socket.")
359 # # return
360 # socketAddr = "/var/ossec/queue/sockets/queue"
361 # from socket import AF_UNIX
362 # from socket import SOCK_DGRAM
363 # from socket import socket
364
365 # if isinstance(msg, str):
366 # try:
367 # msg = json.loads(msg)
368 # except json.JSONDecodeError as e:
369 # logger.error(f"Invalid JSON string: {e}")
370 # raise HTTPException(
371 # status_code=400,
372 # detail="Invalid JSON string.",
373 # )
374 # elif not isinstance(msg, dict):
375 # logger.error("Invalid message type. Expected str or dict.")
376 # raise HTTPException(
377 # status_code=400,
378 # detail="Invalid message type. Expected str or dict.",
379 # )
380
381 # try:
382 # integration = msg["integration"]
383 # except KeyError as e:
384 # logger.error(f"KeyError: {e}")
385 # raise HTTPException(
386 # status_code=400,
387 # detail="Invalid message format. Could not extract 'integration'.",
388 # )
389
390 # socketAddr = "/var/ossec/queue/sockets/queue"
391
392 # try:
393 # msg_str = json.dumps(msg)
394 # logger.info(f"Sending {msg_str} to {socketAddr} socket.")
395 # message = f"1:{integration}:{msg_str}"
396 # sock = socket(AF_UNIX, SOCK_DGRAM)
397 # sock.connect(socketAddr)
398 # sock.send(message.encode())
399 # sock.close()
400 # logger.info(f"Message sent to {socketAddr} socket.")
401 # return {"success": True, "message": "Message sent to Wazuh Socket."}
402 # except Exception as e:
403 # logger.error(f"Error: {e}")
404 # raise HTTPException(
405 # status_code=500,
406 # detail=f"Error: {e}",
407 # )