main
py 590 lines 20.8 KB
Raw
1 import re
2 from typing import Any
3 from typing import Dict
4
5 import httpx
6 from fastapi import HTTPException
7 from loguru import logger
8 from sqlalchemy.ext.asyncio import AsyncSession
9
10 from app.connectors.utils import get_connector_info_from_db
11 from app.db.db_session import get_db_session
12 from app.threat_intel.schema.socfortress import IoCMapping
13 from app.threat_intel.schema.socfortress import IoCResponse
14 from app.threat_intel.schema.socfortress import SocfortressAiAlertRequest
15 from app.threat_intel.schema.socfortress import SocfortressAiAlertResponse
16 from app.threat_intel.schema.socfortress import SocfortressAiWazuhExclusionRuleResponse
17 from app.threat_intel.schema.socfortress import (
18 SocfortressProcessNameAnalysisAPIResponse,
19 )
20 from app.threat_intel.schema.socfortress import SocfortressProcessNameAnalysisRequest
21 from app.threat_intel.schema.socfortress import SocfortressProcessNameAnalysisResponse
22 from app.threat_intel.schema.socfortress import SocfortressThreatIntelRequest
23 from app.threat_intel.schema.socfortress import (
24 VelociraptorArtifactRecommendationRequest,
25 )
26 from app.threat_intel.schema.socfortress import (
27 VelociraptorArtifactRecommendationResponse,
28 )
29 from app.threat_intel.schema.virustotal import VirusTotalResponse
30 from app.utils import get_connector_attribute
31
32
33 async def get_socfortress_threat_intel_attributes(
34 column_name: str,
35 session: AsyncSession,
36 ) -> str:
37 """
38 Gets the SocFortress Threat Intel attribute from the database.
39
40 Args:
41 column_name (str): The column name of the SocFortress Threat Intel attribute.
42 session (AsyncSession): The database session.
43
44 Raises:
45 HTTPException: Raised if the SocFortress Threat Intel Attribute is not found.
46
47 Returns:
48 str: The SocFortress Threat Intel Attribute.
49
50 """
51 attribute_value = await get_connector_attribute(
52 connector_id=10,
53 column_name=column_name,
54 session=session,
55 )
56 # Close the session
57 await session.close()
58 if not attribute_value:
59 raise HTTPException(
60 status_code=500,
61 detail="SocFortress Threat Intel attributes not found in the database.",
62 )
63 return attribute_value
64
65
66 async def verify_socfortress_threat_intel_credentials(
67 attributes: Dict[str, Any],
68 ) -> Dict[str, Any]:
69 """
70 Verifies the SOCFortress Threat Intel credentials.
71
72 Args:
73 attributes (Dict[str, Any]): The connector attributes.
74
75 Returns:
76 Dict[str, Any]: The connector attributes.
77
78 Raises:
79 HTTPException: Raised if the SOCFortress Threat Intel credentials are invalid.
80 """
81 api_key = attributes.get("connector_api_key", None)
82 url = attributes.get("connector_url", None)
83 if api_key is None or url is None:
84 logger.error("No SOCFortress Threat Intel credentials found in the database")
85 raise HTTPException(
86 status_code=500,
87 detail="SOCFortress Threat Intel credentials not found in the database",
88 )
89 return attributes
90
91
92 async def verifiy_socfortress_threat_intel_connector(connector_name: str) -> str:
93 """
94 Verifies the SOCFortress Threat Intel connector.
95
96 Args:
97 connector_name (str): The name of the connector.
98
99 Returns:
100 str: The connector name.
101
102 Raises:
103 HTTPException: Raised if the connector name is not SOCFortress Threat Intel.
104 """
105 logger.info("Verifying SOCFortress Threat Intel connector")
106 async with get_db_session() as session: # This will correctly enter the context manager
107 attributes = await get_connector_info_from_db(connector_name, session)
108 if attributes is None:
109 logger.error("No SOCFortress Threat Intel connector found in the database")
110 return None
111 request = SocfortressThreatIntelRequest(
112 ioc_value="evil.socfortress.co",
113 customer_code="00001",
114 )
115 response = await invoke_socfortress_threat_intel_api(
116 attributes["connector_api_key"],
117 attributes["connector_url"],
118 request,
119 )
120 if "data" in response and response["data"].get("comment") == "This is a test IoC":
121 logger.info("Verified SOCFortress Threat Intel connector")
122 return {
123 "connectionSuccessful": True,
124 "message": "Successfully verified SOCFortress Threat Intel connector",
125 }
126 else:
127 logger.error("Failed to verify SOCFortress Threat Intel connector")
128 return {
129 "connectionSuccessful": False,
130 "message": "Failed to verify SOCFortress Threat Intel connector",
131 }
132
133
134 async def invoke_socfortress_threat_intel_api(
135 api_key: str,
136 url: str,
137 request: SocfortressThreatIntelRequest,
138 ) -> dict:
139 """
140 Invokes the Socfortress Threat Intel API with the provided API key, URL, and request parameters.
141
142 Args:
143 api_key (str): The API key for authentication.
144 url (str): The URL of the Socfortress Threat Intel API.
145 request (SocfortressThreatIntelRequest): The request object containing the IOC value and customer code.
146
147 Returns:
148 dict: The JSON response from the Socfortress Threat Intel API.
149
150 Raises:
151 httpx.HTTPStatusError: If the API request fails with a non-successful status code.
152 """
153 headers = {"module-version": "your_module_version", "x-api-key": api_key}
154 params = {"value": f"{request.ioc_value}&customer_code={request.customer_code}"}
155 logger.info(f"Invoking Socfortress Threat Intel API with params: {params}")
156 async with httpx.AsyncClient() as client:
157 response = await client.get(url, headers=headers, params=params)
158 return response.json()
159
160
161 def determine_ioc_type(ioc_value: str) -> str:
162 """
163 Determine the type of the IOC value and return the appropriate endpoint.
164
165 Args:
166 ioc_value (str): The IOC value.
167
168 Returns:
169 str: The endpoint for the IOC value.
170
171 Raises:
172 ValueError: If the IOC value is invalid.
173 """
174 ip_pattern = re.compile(r"^(?:[0-9]{1,3}\.){3}[0-9]{1,3}$")
175 domain_pattern = re.compile(r"^(?:[a-zA-Z0-9-]+\.)+[a-zA-Z]{2,}$")
176 hash_pattern = re.compile(r"^[a-fA-F0-9]{32}$|^[a-fA-F0-9]{40}$|^[a-fA-F0-9]{64}$")
177 url_pattern = re.compile(r"^(https?|ftp):\/\/[^\s/$.?#].[^\s]*$")
178
179 if ip_pattern.match(ioc_value):
180 return f"/ip_addresses/{ioc_value}"
181 elif domain_pattern.match(ioc_value):
182 return f"/domains/{ioc_value}"
183 elif hash_pattern.match(ioc_value):
184 return f"/files/{ioc_value}"
185 elif url_pattern.match(ioc_value):
186 raise HTTPException(
187 status_code=400,
188 detail="URL scanning is currently not supported.",
189 )
190 return "/urls"
191 else:
192 raise HTTPException(
193 status_code=400,
194 detail="Invalid IOC value provided. Only IP addresses, domains, URLs, and hashes are supported.",
195 )
196
197
198 async def fetch_virustotal_data(api_key: str, full_url: str, ioc_value: str, is_url: bool) -> dict:
199 """
200 Fetch data from the VirusTotal API.
201
202 Args:
203 api_key (str): The API key for authentication.
204 full_url (str): The full URL of the VirusTotal API endpoint.
205 ioc_value (str): The IOC value.
206 is_url (bool): Flag indicating if the IOC value is a URL.
207
208 Returns:
209 dict: The JSON response from the VirusTotal API.
210
211 Raises:
212 httpx.HTTPStatusError: If the API request fails with a non-successful status code.
213 """
214 headers = {"x-apikey": api_key}
215 async with httpx.AsyncClient() as client:
216 if is_url:
217 headers["Content-Type"] = "application/x-www-form-urlencoded"
218 data = {"url": ioc_value}
219 response = await client.post(full_url, headers=headers, data=data)
220 response.raise_for_status()
221 analysis_id = response.json()["data"]["id"]
222 url_report_url = f"https://www.virustotal.com/api/v3/urls/{analysis_id}"
223 response = await client.get(url_report_url, headers=headers)
224 else:
225 response = await client.get(full_url, headers=headers)
226 response.raise_for_status()
227 return VirusTotalResponse.parse_obj(response.json())
228
229
230 async def invoke_virustotal_api(
231 api_key: str,
232 url: str,
233 request: SocfortressThreatIntelRequest,
234 ) -> dict:
235 """
236 Invokes the VirusTotal API with the provided API key, URL, and request parameters.
237
238 Args:
239 api_key (str): The API key for authentication.
240 url (str): The base URL of the VirusTotal API.
241 request (SocfortressThreatIntelRequest): The request object containing the IOC value and customer code.
242
243 Returns:
244 dict: The JSON response from the VirusTotal API.
245
246 Raises:
247 httpx.HTTPStatusError: If the API request fails with a non-successful status code.
248 """
249 ioc_value = request.ioc_value
250 endpoint = determine_ioc_type(ioc_value)
251 full_url = f"{url}{endpoint}"
252 is_url = endpoint == "/urls"
253 return await fetch_virustotal_data(api_key, full_url, ioc_value, is_url)
254
255
256 async def invoke_socfortress_process_name_api(
257 api_key: str,
258 url: str,
259 request: SocfortressProcessNameAnalysisRequest,
260 ) -> dict:
261 """
262 Invokes the Socfortress Process Analysis API with the provided API key, URL, and request parameters.
263
264 Args:
265 api_key (str): The API key for authentication.
266 url (str): The URL of the Socfortress Intel URL
267 request (SocfortressProcessNameAnalysisRequest): The request object containing the Process Name
268
269 Returns:
270 dict: The JSON response from the Process Name Analysis API.
271
272 Raises:
273 httpx.HTTPStatusError: If the API request fails with a non-successful status code.
274 """
275 headers = {"module-version": "your_module_version", "x-api-key": api_key}
276 params = {"value": f"{request.process_name}"}
277 logger.info(f"Invoking Socfortress Process Name Analysis with params: {params} and headers: {headers} and url: {url}")
278 async with httpx.AsyncClient() as client:
279 response = await client.get(url, headers=headers, params=params)
280 return response.json()
281
282
283 async def invoke_socfortress_ai_alert_api(
284 api_key: str,
285 url: str,
286 request: SocfortressAiAlertRequest,
287 timeout: int = 60,
288 ) -> dict:
289 """
290 Invokes the Socfortress AI Alert API with the provided API key, URL, and request parameters.
291
292 Args:
293 api_key (str): The API key for authentication.
294 url (str): The URL of the Socfortress Intel URL
295 request (SocfortressAiAlertRequest): The request object containing the Process Name
296
297 Returns:
298 dict: The JSON response from the AI Alert API.
299
300 Raises:
301 httpx.HTTPStatusError: If the API request fails with a non-successful status code.
302 """
303 headers = {"module-version": "1.0", "x-api-key": api_key}
304 try:
305 async with httpx.AsyncClient(timeout=timeout) as client:
306 response = await client.post(url, json=request.model_dump(), headers=headers)
307 response.raise_for_status() # Raise an exception for non-successful status codes
308 return response.json()
309 except httpx.HTTPStatusError as e:
310 if e.response.status_code == 429:
311 logger.error(f"Rate limit reached: {e.response.status_code} - {e.response.text}")
312 raise HTTPException(
313 status_code=429,
314 detail="Rate limit reached for the month. Please try again next month.",
315 )
316 else:
317 logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
318 raise HTTPException(
319 status_code=e.response.status_code,
320 detail=f"HTTP error occurred: {e.response.status_code} - {e.response.text}",
321 )
322 except httpx.RequestError as e:
323 logger.error(f"Request error occurred: {e}")
324 raise HTTPException(
325 status_code=500,
326 detail=f"Request error occurred: {e}",
327 )
328 except Exception as e:
329 logger.error(f"An unexpected error occurred: {e}")
330 raise HTTPException(
331 status_code=500,
332 detail=f"An unexpected error occurred: {e}",
333 )
334
335
336 async def get_ioc_response(
337 license_key: str,
338 request: SocfortressThreatIntelRequest,
339 session: AsyncSession,
340 ) -> IoCResponse:
341 """
342 Retrieves IoC response from Socfortress Threat Intel API.
343
344 Args:
345 request (SocfortressThreatIntelRequest): The request object containing the IoC data.
346 session (AsyncSession): The async session object for making HTTP requests.
347
348 Returns:
349 IoCResponse: The response object containing the IoC data and success status.
350 """
351 url = "https://intel.socfortress.co/search"
352 response_data = await invoke_socfortress_threat_intel_api(license_key, url, request)
353
354 # Using .get() with default values
355 data = response_data.get("data", {})
356 success = response_data.get("success", False)
357 message = response_data.get("message", "No message provided")
358
359 return IoCResponse(data=IoCMapping(**data), success=success, message=message)
360
361
362 async def get_ai_alert_response(
363 license_key: str,
364 request: SocfortressAiAlertRequest,
365 ) -> SocfortressAiAlertResponse:
366 """
367 Retrieves IoC response from Socfortress Threat Intel API.
368
369 Args:
370 request (SocfortressAiAlertRequest): The request object containing the IoC data.
371 session (AsyncSession): The async session object for making HTTP requests.
372
373 Returns:
374 SocfortressAiAlertResponse: The response object containing the IoC data and success status.
375 """
376 url = "https://ai.socfortress.co/analyze-alert"
377
378 response_data = await invoke_socfortress_ai_alert_api(license_key, url, request)
379
380 # If message is `Forbidden`, raise an HTTPException
381 if response_data.get("message") == "Forbidden":
382 raise HTTPException(
383 status_code=403,
384 detail="Forbidden access to the Socfortress AI Alert API",
385 )
386
387 return SocfortressAiAlertResponse(**response_data)
388
389
390 async def get_wazuh_exclusion_rule_response(
391 license_key: str,
392 request: SocfortressAiAlertRequest,
393 ) -> SocfortressAiWazuhExclusionRuleResponse:
394 """
395 Retrieves IoC response from Socfortress Threat Intel API.
396
397 Args:
398 request (SocfortressAiAlertRequest): The request object containing the IoC data.
399 session (AsyncSession): The async session object for making HTTP requests.
400
401 Returns:
402 SocfortressAiWazuhExclusionRuleResponse: The response object containing the IoC data and success status.
403 """
404 url = "https://ai.socfortress.co/wazuh-exclusion-rule"
405
406 response_data = await invoke_socfortress_ai_alert_api(license_key, url, request)
407
408 # If message is `Forbidden`, raise an HTTPException
409 if response_data.get("message") == "Forbidden":
410 raise HTTPException(
411 status_code=403,
412 detail="Forbidden access to the Socfortress AI Alert API",
413 )
414
415 return SocfortressAiWazuhExclusionRuleResponse(**response_data)
416
417
418 async def get_velociraptor_artifact_recommendation_response(
419 license_key: str,
420 request: VelociraptorArtifactRecommendationRequest,
421 ) -> VelociraptorArtifactRecommendationResponse:
422 """
423 Retrieves Artifact recommendation response from Socfortress Threat Intel API.
424
425 Args:
426 request (VelociraptorArtifactRecommendationRequest): The request object containing the alert data.
427 session (AsyncSession): The async session object for making HTTP requests.
428
429 Returns:
430 VelociraptorArtifactRecommendationResponse: The response object containing the artifact recommendation data and success status.
431 """
432 url = "https://ai.socfortress.co/velociraptor-artifact-recommendation"
433
434 response_data = await invoke_socfortress_ai_alert_api(license_key, url, request)
435
436 # If message is `Forbidden`, raise an HTTPException
437 if response_data.get("message") == "Forbidden":
438 raise HTTPException(
439 status_code=403,
440 detail="Forbidden access to the Socfortress AI Alert API",
441 )
442 elif "429" in response_data.get("detail", ""):
443 raise HTTPException(
444 status_code=429,
445 detail="Message is too large. Please try again with a smaller message.",
446 )
447 elif "too large" in response_data.get("detail", ""):
448 raise HTTPException(
449 status_code=429,
450 detail="Message is too large. Please try again with a smaller message.",
451 )
452
453 return VelociraptorArtifactRecommendationResponse(**response_data)
454
455
456 async def get_process_analysis_response(
457 license_key: str,
458 request: SocfortressProcessNameAnalysisRequest,
459 session: AsyncSession,
460 ) -> SocfortressProcessNameAnalysisResponse:
461 """
462 Retrieves IoC response from Socfortress Threat Intel API.
463
464 Args:
465 request (SocfortressProcessNameAnalysisRequest): The request object containing the IoC data.
466 session (AsyncSession): The async session object for making HTTP requests.
467
468 Returns:
469 SocfortressProcessNameAnalysisResponse: The response object containing the IoC data and success status.
470 """
471 url = "https://processname.socfortress.co/search"
472 response_data = await invoke_socfortress_process_name_api(license_key, url, request)
473
474 # If message is `Forbidden`, raise an HTTPException
475 if response_data.get("message") == "Forbidden":
476 raise HTTPException(
477 status_code=403,
478 detail="Forbidden access to the Socfortress Process Name Analysis API",
479 )
480
481 # Using .get() with default values
482 data = response_data.get("data", {})
483 success = response_data.get("success", False)
484 message = response_data.get("message", "No message provided")
485
486 return SocfortressProcessNameAnalysisResponse(data=SocfortressProcessNameAnalysisAPIResponse(**data), success=success, message=message)
487
488
489 async def socfortress_threat_intel_lookup(
490 lincense_key: str,
491 request: SocfortressThreatIntelRequest,
492 session: AsyncSession,
493 ) -> SocfortressProcessNameAnalysisResponse:
494 """
495 Performs a threat intelligence lookup using the Socfortress service.
496
497 Args:
498 request (SocfortressThreatIntelRequest): The request object containing the IoC to lookup.
499 session (AsyncSession): The async session object for making HTTP requests.
500
501 Returns:
502 IoCResponse: The response object containing the threat intelligence information.
503 """
504 return await get_ioc_response(
505 license_key=lincense_key,
506 request=request,
507 session=session,
508 )
509
510
511 async def socfortress_process_analysis_lookup(
512 lincense_key: str,
513 request: SocfortressProcessNameAnalysisRequest,
514 session: AsyncSession,
515 ) -> IoCResponse:
516 """
517 Performs a process analysis intelligence lookup using the Socfortress service.
518
519 Args:
520 request (SocfortressThreatIntelRequest): The request object containing the IoC to lookup.
521 session (AsyncSession): The async session object for making HTTP requests.
522
523 Returns:
524 IoCResponse: The response object containing the threat intelligence information.
525 """
526 return await get_process_analysis_response(
527 license_key=lincense_key,
528 request=request,
529 session=session,
530 )
531
532
533 async def socfortress_ai_alert_lookup(
534 lincense_key: str,
535 request: SocfortressAiAlertRequest,
536 ) -> SocfortressAiAlertResponse:
537 """
538 Performs a AI alert lookup using the Socfortress service.
539
540 Args:
541 request (SocfortressAiAlertRequest): The request object containing the IoC to lookup.
542 session (AsyncSession): The async session object for making HTTP requests.
543
544 Returns:
545 IoCResponse: The response object containing the threat intelligence information.
546 """
547 return await get_ai_alert_response(
548 license_key=lincense_key,
549 request=request,
550 )
551
552
553 async def socfortress_wazuh_exclusion_rule_lookup(
554 lincense_key: str,
555 request: SocfortressAiAlertRequest,
556 ) -> SocfortressAiWazuhExclusionRuleResponse:
557 """
558 Performs a AI alert lookup using the Socfortress service.
559
560 Args:
561 request (SocfortressAiAlertRequest): The request object containing the IoC to lookup.
562 session (AsyncSession): The async session object for making HTTP requests.
563
564 Returns:
565 IoCResponse: The response object containing the threat intelligence information.
566 """
567 return await get_wazuh_exclusion_rule_response(
568 license_key=lincense_key,
569 request=request,
570 )
571
572
573 async def socfortress_velociraptor_recommendation_lookup(
574 lincense_key: str,
575 request: VelociraptorArtifactRecommendationRequest,
576 ) -> VelociraptorArtifactRecommendationResponse:
577 """
578 Performs a AI alert lookup using the Socfortress service.
579
580 Args:
581 request (VelociraptorArtifactRecommendationRequest): The request object containing the IoC to lookup.
582 session (AsyncSession): The async session object for making HTTP requests.
583
584 Returns:
585 IoCResponse: The response object containing the threat intelligence information.
586 """
587 return await get_velociraptor_artifact_recommendation_response(
588 license_key=lincense_key,
589 request=request,
590 )