| 1 | from datetime import datetime |
| 2 | |
| 3 | from fastapi import APIRouter |
| 4 | from fastapi import Depends |
| 5 | from fastapi import File |
| 6 | from fastapi import Form |
| 7 | from fastapi import HTTPException |
| 8 | from fastapi import Security |
| 9 | from fastapi import UploadFile |
| 10 | from loguru import logger |
| 11 | from sqlalchemy.ext.asyncio import AsyncSession |
| 12 | |
| 13 | from app.agents.services.status import get_agent_os_by_id |
| 14 | from app.auth.utils import AuthHandler |
| 15 | from app.connectors.velociraptor.services.artifacts import get_artifacts |
| 16 | from app.db.db_session import get_db |
| 17 | from app.incidents.schema.db_operations import CommentCreate |
| 18 | from app.incidents.schema.incident_alert import CreateAlertRequest |
| 19 | from app.incidents.schema.incident_alert import CreateAlertRequestRoute |
| 20 | from app.incidents.schema.incident_alert import GenericAlertModel |
| 21 | from app.incidents.services.db_operations import create_comment |
| 22 | from app.incidents.services.incident_alert import get_single_alert_details |
| 23 | from app.integrations.copilot_mcp.routes.copilot_mcp import query_mcp |
| 24 | from app.integrations.copilot_mcp.schema.copilot_mcp import MCPQueryRequest |
| 25 | from app.integrations.copilot_mcp.schema.copilot_mcp import MCPQueryResponse |
| 26 | from app.middleware.license import get_license |
| 27 | from app.middleware.license import is_feature_enabled |
| 28 | |
| 29 | # from app.threat_intel.schema.socfortress import SocfortressProcessNameAnalysisResponse |
| 30 | from app.threat_intel.schema.socfortress import IoCResponse |
| 31 | from app.threat_intel.schema.socfortress import SocfortressAiAlertRequest |
| 32 | from app.threat_intel.schema.socfortress import SocfortressAiAlertResponse |
| 33 | from app.threat_intel.schema.socfortress import SocfortressAiWazuhExclusionRuleResponse |
| 34 | from app.threat_intel.schema.socfortress import SocfortressProcessNameAnalysisRequest |
| 35 | from app.threat_intel.schema.socfortress import SocfortressThreatIntelRequest |
| 36 | from app.threat_intel.schema.socfortress import ( |
| 37 | VelociraptorArtifactRecommendationRequest, |
| 38 | ) |
| 39 | from app.threat_intel.schema.socfortress import ( |
| 40 | VelociraptorArtifactRecommendationResponse, |
| 41 | ) |
| 42 | from app.threat_intel.schema.socfortress import VirusTotalThreatIntelRequest |
| 43 | from app.threat_intel.schema.virustotal import FileAnalysisResponse |
| 44 | from app.threat_intel.schema.virustotal import FileReportResponse |
| 45 | from app.threat_intel.schema.virustotal import FileSubmissionRequest |
| 46 | from app.threat_intel.schema.virustotal import FileSubmissionResponse |
| 47 | from app.threat_intel.schema.virustotal import VirusTotalRouteResponse |
| 48 | |
| 49 | # from app.threat_intel.services.socfortress import socfortress_process_analysis_lookup |
| 50 | from app.threat_intel.services.socfortress import invoke_virustotal_api |
| 51 | from app.threat_intel.services.socfortress import socfortress_ai_alert_lookup |
| 52 | from app.threat_intel.services.socfortress import socfortress_threat_intel_lookup |
| 53 | from app.threat_intel.services.socfortress import ( |
| 54 | socfortress_velociraptor_recommendation_lookup, |
| 55 | ) |
| 56 | from app.threat_intel.services.socfortress import ( |
| 57 | socfortress_wazuh_exclusion_rule_lookup, |
| 58 | ) |
| 59 | from app.threat_intel.services.virustotal_file import get_file_analysis_status |
| 60 | from app.threat_intel.services.virustotal_file import get_file_report |
| 61 | from app.threat_intel.services.virustotal_file import submit_and_wait_for_analysis |
| 62 | from app.threat_intel.services.virustotal_file import submit_file_to_virustotal |
| 63 | from app.utils import get_connector_attribute |
| 64 | |
| 65 | # App specific imports |
| 66 | |
| 67 | threat_intel_socfortress_router = APIRouter() |
| 68 | |
| 69 | |
| 70 | async def ensure_api_key_exists(session: AsyncSession = Depends(get_db)) -> bool: |
| 71 | """ |
| 72 | Ensures that the SocFortress API key exists in the database. |
| 73 | |
| 74 | Args: |
| 75 | session (AsyncSession): The database session. |
| 76 | |
| 77 | Raises: |
| 78 | HTTPException: Raised if the SocFortress API key is not found. |
| 79 | |
| 80 | Returns: |
| 81 | bool: True if the API key exists, otherwise raises HTTPException. |
| 82 | """ |
| 83 | api_key = await get_connector_attribute( |
| 84 | connector_id=10, |
| 85 | column_name="connector_api_key", |
| 86 | session=session, |
| 87 | ) |
| 88 | # Close the session |
| 89 | await session.close() |
| 90 | if not api_key: |
| 91 | raise HTTPException( |
| 92 | status_code=500, |
| 93 | detail="SocFortress API key not found in the database.", |
| 94 | ) |
| 95 | return True |
| 96 | |
| 97 | |
| 98 | async def ensure_virustotal_connector(session: AsyncSession = Depends(get_db)) -> dict: |
| 99 | """ |
| 100 | Ensures that the VirusTotal connector is properly configured. |
| 101 | |
| 102 | Args: |
| 103 | session (AsyncSession): The database session dependency |
| 104 | |
| 105 | Returns: |
| 106 | dict: Dictionary containing API key and URL |
| 107 | |
| 108 | Raises: |
| 109 | HTTPException: If connector is not configured or verified |
| 110 | """ |
| 111 | # Check if the connector is verified |
| 112 | if not await get_connector_attribute( |
| 113 | connector_name="VirusTotal", |
| 114 | column_name="connector_verified", |
| 115 | session=session, |
| 116 | ): |
| 117 | raise HTTPException( |
| 118 | status_code=500, |
| 119 | detail="VirusTotal connector is not verified.", |
| 120 | ) |
| 121 | |
| 122 | api_key = await get_connector_attribute( |
| 123 | connector_name="VirusTotal", |
| 124 | column_name="connector_api_key", |
| 125 | session=session, |
| 126 | ) |
| 127 | |
| 128 | url = await get_connector_attribute( |
| 129 | connector_name="VirusTotal", |
| 130 | column_name="connector_url", |
| 131 | session=session, |
| 132 | ) |
| 133 | |
| 134 | if not api_key: |
| 135 | raise HTTPException( |
| 136 | status_code=500, |
| 137 | detail="VirusTotal API key not found in the database.", |
| 138 | ) |
| 139 | |
| 140 | return {"api_key": api_key, "url": url} |
| 141 | |
| 142 | |
| 143 | @threat_intel_socfortress_router.post( |
| 144 | "/socfortress", |
| 145 | response_model=IoCResponse, |
| 146 | description="SocFortress Threat Intel", |
| 147 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 148 | ) |
| 149 | async def threat_intel_socfortress( |
| 150 | request: SocfortressThreatIntelRequest, |
| 151 | session: AsyncSession = Depends(get_db), |
| 152 | # _key_exists: bool = Depends(ensure_api_key_exists), |
| 153 | ): |
| 154 | """ |
| 155 | Endpoint for SocFortress Threat Intel. |
| 156 | |
| 157 | This endpoint allows authorized users with 'admin' or 'analyst' scope to perform SocFortress threat intelligence lookup. |
| 158 | |
| 159 | Parameters: |
| 160 | - request: SocfortressThreatIntelRequest - The request payload containing the necessary information for the lookup. |
| 161 | - session: AsyncSession (optional) - The database session to use for the lookup. |
| 162 | - _key_exists: bool (optional) - A dependency to ensure the API key exists. |
| 163 | |
| 164 | Returns: |
| 165 | - IoCResponse: The response model containing the results of the SocFortress threat intelligence lookup. |
| 166 | """ |
| 167 | await is_feature_enabled("THREAT INTEL", session=session) |
| 168 | logger.info("Running SOCFortress Threat Intel. Grabbing License") |
| 169 | |
| 170 | socfortress_lookup = await socfortress_threat_intel_lookup( |
| 171 | lincense_key=(await get_license(session)).license_key, |
| 172 | request=request, |
| 173 | session=session, |
| 174 | ) |
| 175 | return socfortress_lookup |
| 176 | |
| 177 | |
| 178 | @threat_intel_socfortress_router.post( |
| 179 | "/virustotal", |
| 180 | response_model=VirusTotalRouteResponse, |
| 181 | description="VirusTotal Enrichment Threat Intel", |
| 182 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 183 | ) |
| 184 | async def threat_intel_virustotal( |
| 185 | request: VirusTotalThreatIntelRequest, |
| 186 | session: AsyncSession = Depends(get_db), |
| 187 | ): |
| 188 | """ |
| 189 | Endpoint for VirusTotal Threat Intel. |
| 190 | |
| 191 | This endpoint allows authorized users with 'admin' or 'analyst' scope to perform VirusTotal threat intelligence lookup. |
| 192 | |
| 193 | Parameters: |
| 194 | - request: VirusTotalThreatIntelRequest - The request payload containing the necessary information for the lookup. |
| 195 | - session: AsyncSession (optional) - The database session to use for the lookup. |
| 196 | - _key_exists: bool (optional) - A dependency to ensure the API key exists. |
| 197 | |
| 198 | Returns: |
| 199 | - IoCResponse: The response model containing the results of the VirusTotal threat intelligence lookup. |
| 200 | """ |
| 201 | logger.info("Running VirusTotal Threat Intel.") |
| 202 | |
| 203 | # Check if the connector is verified |
| 204 | if not await get_connector_attribute( |
| 205 | connector_name="VirusTotal", |
| 206 | column_name="connector_verified", |
| 207 | session=session, |
| 208 | ): |
| 209 | raise HTTPException( |
| 210 | status_code=500, |
| 211 | detail="VirusTotal connector is not verified.", |
| 212 | ) |
| 213 | |
| 214 | return VirusTotalRouteResponse( |
| 215 | data=await invoke_virustotal_api( |
| 216 | url=await get_connector_attribute( |
| 217 | connector_name="VirusTotal", |
| 218 | column_name="connector_url", |
| 219 | session=session, |
| 220 | ), |
| 221 | api_key=await get_connector_attribute( |
| 222 | connector_name="VirusTotal", |
| 223 | column_name="connector_api_key", |
| 224 | session=session, |
| 225 | ), |
| 226 | request=request, |
| 227 | ), |
| 228 | success=True, |
| 229 | message="VirusTotal threat intelligence lookup was successful.", |
| 230 | ) |
| 231 | |
| 232 | |
| 233 | @threat_intel_socfortress_router.post( |
| 234 | "/virustotal/file/submit", |
| 235 | response_model=FileSubmissionResponse, |
| 236 | description="Submit a file to VirusTotal for analysis", |
| 237 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 238 | ) |
| 239 | async def submit_file_for_analysis( |
| 240 | file: UploadFile = File(..., description="File to analyze (max 32MB for free API)"), |
| 241 | password: str = Form(None, description="Password for encrypted files"), |
| 242 | vt_config: dict = Depends(ensure_virustotal_connector), |
| 243 | ): |
| 244 | """ |
| 245 | Submit a file to VirusTotal for malware analysis. |
| 246 | |
| 247 | This endpoint allows authorized users to upload files for analysis. |
| 248 | The file will be submitted to VirusTotal and an analysis ID will be returned. |
| 249 | |
| 250 | Parameters: |
| 251 | - file: UploadFile - The file to be analyzed |
| 252 | - password: str (optional) - Password for encrypted files |
| 253 | - vt_config: dict - VirusTotal connector configuration (injected dependency) |
| 254 | |
| 255 | Returns: |
| 256 | - FileSubmissionResponse: Contains the analysis ID for tracking |
| 257 | """ |
| 258 | logger.info(f"Submitting file {file.filename} to VirusTotal for analysis") |
| 259 | |
| 260 | # Validate file size (32MB limit for free API) |
| 261 | if file.size and file.size > 32 * 1024 * 1024: # 32MB |
| 262 | raise HTTPException(status_code=413, detail="File too large. Maximum file size is 32MB for free API keys.") |
| 263 | |
| 264 | # Validate password is only provided for zip files |
| 265 | if password: |
| 266 | # Check file extension |
| 267 | filename = file.filename or "" |
| 268 | file_extension = filename.lower().split(".")[-1] if "." in filename else "" |
| 269 | |
| 270 | # Check MIME type |
| 271 | content_type = file.content_type or "" |
| 272 | |
| 273 | # Define valid zip file indicators |
| 274 | zip_extensions = ["zip", "zipx"] |
| 275 | zip_mime_types = ["application/zip", "application/x-zip-compressed", "application/x-zip", "multipart/x-zip"] |
| 276 | |
| 277 | # Validate that it's a zip file |
| 278 | is_zip_extension = file_extension in zip_extensions |
| 279 | is_zip_mime = content_type in zip_mime_types |
| 280 | |
| 281 | if not (is_zip_extension or is_zip_mime): |
| 282 | raise HTTPException( |
| 283 | status_code=400, |
| 284 | detail=f"Password can only be provided for zip files. " |
| 285 | f"File '{filename}' has extension '{file_extension}' and MIME type '{content_type}', " |
| 286 | f"which are not recognized as zip file formats.", |
| 287 | ) |
| 288 | |
| 289 | # Create request object |
| 290 | request = FileSubmissionRequest(password=password) |
| 291 | |
| 292 | # Submit the file |
| 293 | return await submit_file_to_virustotal(api_key=vt_config["api_key"], file=file, request=request) |
| 294 | |
| 295 | |
| 296 | @threat_intel_socfortress_router.get( |
| 297 | "/virustotal/analysis/{analysis_id}", |
| 298 | response_model=FileAnalysisResponse, |
| 299 | description="Get the status of a file analysis", |
| 300 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 301 | ) |
| 302 | async def get_analysis_status( |
| 303 | analysis_id: str, |
| 304 | vt_config: dict = Depends(ensure_virustotal_connector), |
| 305 | ): |
| 306 | """ |
| 307 | Get the current status of a file analysis. |
| 308 | |
| 309 | Parameters: |
| 310 | - analysis_id: str - The analysis ID returned from file submission |
| 311 | - vt_config: dict - VirusTotal connector configuration (injected dependency) |
| 312 | |
| 313 | Returns: |
| 314 | - FileAnalysisResponse: Current analysis status and results |
| 315 | """ |
| 316 | logger.info(f"Getting analysis status for ID: {analysis_id}") |
| 317 | |
| 318 | return await get_file_analysis_status(api_key=vt_config["api_key"], analysis_id=analysis_id) |
| 319 | |
| 320 | |
| 321 | @threat_intel_socfortress_router.get( |
| 322 | "/virustotal/file/{file_id}", |
| 323 | response_model=FileReportResponse, |
| 324 | description="Get detailed analysis report for a file", |
| 325 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 326 | ) |
| 327 | async def get_file_analysis_report( |
| 328 | file_id: str, |
| 329 | vt_config: dict = Depends(ensure_virustotal_connector), |
| 330 | ): |
| 331 | """ |
| 332 | Get the detailed analysis report for a file. |
| 333 | |
| 334 | Parameters: |
| 335 | - file_id: str - The file ID (hash) to get the report for |
| 336 | - vt_config: dict - VirusTotal connector configuration (injected dependency) |
| 337 | |
| 338 | Returns: |
| 339 | - FileReportResponse: Detailed analysis report |
| 340 | """ |
| 341 | logger.info(f"Getting file report for ID: {file_id}") |
| 342 | |
| 343 | return await get_file_report(api_key=vt_config["api_key"], file_id=file_id) |
| 344 | |
| 345 | |
| 346 | @threat_intel_socfortress_router.post( |
| 347 | "/virustotal/file/analyze", |
| 348 | response_model=FileReportResponse, |
| 349 | description="Submit a file and wait for analysis completion", |
| 350 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 351 | ) |
| 352 | async def analyze_file_complete( |
| 353 | file: UploadFile = File(..., description="File to analyze (max 32MB for free API)"), |
| 354 | password: str = Form(None, description="Password for encrypted files"), |
| 355 | max_wait_time: int = Form(300, description="Maximum wait time in seconds (default: 300)"), |
| 356 | poll_interval: int = Form(10, description="Polling interval in seconds (default: 10)"), |
| 357 | vt_config: dict = Depends(ensure_virustotal_connector), |
| 358 | ): |
| 359 | """ |
| 360 | Submit a file to VirusTotal and wait for the analysis to complete. |
| 361 | |
| 362 | This endpoint combines file submission and result retrieval into a single call. |
| 363 | It will wait for the analysis to complete before returning the results. |
| 364 | |
| 365 | Parameters: |
| 366 | - file: UploadFile - The file to be analyzed |
| 367 | - password: str (optional) - Password for encrypted files |
| 368 | - max_wait_time: int - Maximum time to wait for analysis completion (seconds) |
| 369 | - poll_interval: int - Time between status checks (seconds) |
| 370 | - vt_config: dict - VirusTotal connector configuration (injected dependency) |
| 371 | |
| 372 | Returns: |
| 373 | - FileReportResponse: Complete analysis report |
| 374 | """ |
| 375 | logger.info(f"Starting complete analysis for file {file.filename}") |
| 376 | |
| 377 | # Validate file size |
| 378 | if file.size and file.size > 32 * 1024 * 1024: # 32MB |
| 379 | raise HTTPException(status_code=413, detail="File too large. Maximum file size is 32MB for free API keys.") |
| 380 | |
| 381 | # Validate wait time parameters |
| 382 | if max_wait_time < 30 or max_wait_time > 600: # 30 seconds to 10 minutes |
| 383 | raise HTTPException(status_code=400, detail="max_wait_time must be between 30 and 600 seconds") |
| 384 | |
| 385 | if poll_interval < 5 or poll_interval > 60: # 5 seconds to 1 minute |
| 386 | raise HTTPException(status_code=400, detail="poll_interval must be between 5 and 60 seconds") |
| 387 | |
| 388 | # Create request object |
| 389 | request = FileSubmissionRequest(password=password) |
| 390 | |
| 391 | # Submit and wait for analysis |
| 392 | return await submit_and_wait_for_analysis( |
| 393 | api_key=vt_config["api_key"], |
| 394 | file=file, |
| 395 | request=request, |
| 396 | max_wait_time=max_wait_time, |
| 397 | poll_interval=poll_interval, |
| 398 | ) |
| 399 | |
| 400 | |
| 401 | @threat_intel_socfortress_router.post( |
| 402 | "/process_name", |
| 403 | response_model=MCPQueryResponse, |
| 404 | description="SocFortress Process Name Evaluation", |
| 405 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 406 | ) |
| 407 | async def process_name_intel_socfortress( |
| 408 | request: SocfortressProcessNameAnalysisRequest, |
| 409 | session: AsyncSession = Depends(get_db), |
| 410 | ): |
| 411 | """ |
| 412 | Endpoint for SocFortress Process Name Evaluation. |
| 413 | |
| 414 | This endpoint allows authorized users with 'admin' or 'analyst' scope to perform SocFortress process name evaluation. |
| 415 | |
| 416 | Parameters: |
| 417 | - request: SocfortressThreatIntelRequest - The request payload containing the necessary information for the lookup. |
| 418 | - session: AsyncSession (optional) - The database session to use for the lookup. |
| 419 | - _key_exists: bool (optional) - A dependency to ensure the API key exists. |
| 420 | |
| 421 | Returns: |
| 422 | - MCPQueryResponse: The response model containing the results of the SocFortress process name analysis lookup. |
| 423 | """ |
| 424 | # await is_feature_enabled("PROCESS ANALYSIS", session=session) |
| 425 | logger.info("Running SOCFortress Process Name Analysis. Grabbing License") |
| 426 | |
| 427 | return await query_mcp( |
| 428 | MCPQueryRequest( |
| 429 | mcp_server="cyber-news", |
| 430 | input=f"Analyze the process name: {request.process_name}. Provide a risk assessment and any relevant details.", |
| 431 | ), |
| 432 | session=session, |
| 433 | ) |
| 434 | |
| 435 | |
| 436 | async def current_time(): |
| 437 | return datetime.now().strftime("%Y-%m-%d %H:%M:%S") |
| 438 | |
| 439 | |
| 440 | @threat_intel_socfortress_router.post( |
| 441 | "/ai/analyze-alert", |
| 442 | response_model=SocfortressAiAlertResponse, |
| 443 | description="SocFortress Process Name Evaluation", |
| 444 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 445 | ) |
| 446 | async def ai_anaylze_alert_socfortress( |
| 447 | request: CreateAlertRequestRoute, |
| 448 | session: AsyncSession = Depends(get_db), |
| 449 | ): |
| 450 | # Fetch alert details |
| 451 | alert_details = await get_single_alert_details(CreateAlertRequest(index_name=request.index_name, alert_id=request.index_id)) |
| 452 | |
| 453 | assert isinstance(alert_details, GenericAlertModel) |
| 454 | |
| 455 | ai_request = SocfortressAiAlertRequest( |
| 456 | integration="SOCFORTRESS AI", |
| 457 | alert_payload=alert_details.source.model_dump(), |
| 458 | ) |
| 459 | |
| 460 | socfortress_lookup = await socfortress_ai_alert_lookup( |
| 461 | lincense_key=(await get_license(session)).license_key, |
| 462 | request=ai_request, |
| 463 | ) |
| 464 | |
| 465 | await create_comment( |
| 466 | CommentCreate( |
| 467 | alert_id=request.alert_id, |
| 468 | comment=f"SOCFortress AI Analysis: {socfortress_lookup.analysis}", |
| 469 | user_name="admin", |
| 470 | created_at=datetime.now(), |
| 471 | ), |
| 472 | db=session, |
| 473 | ) |
| 474 | |
| 475 | return socfortress_lookup |
| 476 | |
| 477 | |
| 478 | @threat_intel_socfortress_router.post( |
| 479 | "/ai/wazuh-exclusion-rule", |
| 480 | response_model=SocfortressAiWazuhExclusionRuleResponse, |
| 481 | description="SocFortress Process Name Evaluation", |
| 482 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 483 | ) |
| 484 | async def ai_wazuh_exclusion_rule_socfortress( |
| 485 | request: CreateAlertRequestRoute, |
| 486 | session: AsyncSession = Depends(get_db), |
| 487 | ): |
| 488 | # Fetch alert details |
| 489 | alert_details = await get_single_alert_details(CreateAlertRequest(index_name=request.index_name, alert_id=request.index_id)) |
| 490 | |
| 491 | assert isinstance(alert_details, GenericAlertModel) |
| 492 | |
| 493 | ai_request = SocfortressAiAlertRequest( |
| 494 | integration="SOCFORTRESS AI", |
| 495 | alert_payload=alert_details.source.model_dump(), |
| 496 | ) |
| 497 | |
| 498 | logger.info(f"Sending request: {request}") |
| 499 | |
| 500 | socfortress_lookup = await socfortress_wazuh_exclusion_rule_lookup( |
| 501 | lincense_key=(await get_license(session)).license_key, |
| 502 | request=ai_request, |
| 503 | ) |
| 504 | |
| 505 | await create_comment( |
| 506 | CommentCreate( |
| 507 | alert_id=request.alert_id, |
| 508 | comment=f"SOCFortress AI Analysis: {socfortress_lookup.wazuh_exclusion_rule}/n/n{socfortress_lookup.wazuh_exclusion_rule_justification}", |
| 509 | user_name="admin", |
| 510 | created_at=datetime.now(), |
| 511 | ), |
| 512 | db=session, |
| 513 | ) |
| 514 | return socfortress_lookup |
| 515 | |
| 516 | |
| 517 | async def fetch_agent_os(agent_id: str, session: AsyncSession) -> str: |
| 518 | """ |
| 519 | Fetch the operating system of the agent. |
| 520 | |
| 521 | Args: |
| 522 | agent_id (str): The ID of the agent. |
| 523 | session (AsyncSession): The database session. |
| 524 | |
| 525 | Returns: |
| 526 | str: The normalized operating system name. |
| 527 | |
| 528 | Raises: |
| 529 | HTTPException: If the agent OS is not found or unsupported. |
| 530 | """ |
| 531 | agent_os = await get_agent_os_by_id(agent_id=agent_id, session=session) |
| 532 | |
| 533 | if agent_os is None: |
| 534 | raise HTTPException( |
| 535 | status_code=404, |
| 536 | detail="Agent OS not found.", |
| 537 | ) |
| 538 | |
| 539 | # Normalize the OS name |
| 540 | agent_os_lower = agent_os.lower() |
| 541 | if "windows" in agent_os_lower: |
| 542 | return "Windows" |
| 543 | elif "linux" in agent_os_lower or "ubuntu" in agent_os_lower: |
| 544 | return "Linux" |
| 545 | elif "macos" in agent_os_lower or "mac" in agent_os_lower: |
| 546 | return "MacOS" |
| 547 | else: |
| 548 | raise HTTPException( |
| 549 | status_code=400, |
| 550 | detail="Unsupported OS type.", |
| 551 | ) |
| 552 | |
| 553 | |
| 554 | async def filter_artifacts_by_os(artifacts, os): |
| 555 | # Only get the artifacts that start with `Windows.`, `Linux.`, or `MacOS.` |
| 556 | os_artifacts = ["Windows", "Linux", "MacOS"] |
| 557 | os_artifacts = [os_artifact for os_artifact in os_artifacts if os_artifact in os] |
| 558 | |
| 559 | # Artifacts to be stripped out |
| 560 | excluded_artifacts = { |
| 561 | "Windows.Sysinternals.SysmonInstall", |
| 562 | "Windows.Sysinternals.SysmonLogForward", |
| 563 | "Windows.Sysinternals.Autoruns", |
| 564 | "Windows.Sigma.EventLogs", |
| 565 | "Windows.Remediation.Quarantine", |
| 566 | "Windows.Remediation.QuarantineMonitor", |
| 567 | "Windows.Custom.InstallHuntress", |
| 568 | "Windows.Applications.TeamViewer.Incoming", |
| 569 | } |
| 570 | |
| 571 | return [ |
| 572 | artifact |
| 573 | for artifact in artifacts |
| 574 | if any(artifact.name.startswith(os_artifact + ".") for os_artifact in os_artifacts) and artifact.name not in excluded_artifacts |
| 575 | ] |
| 576 | |
| 577 | |
| 578 | async def fetch_artifacts(os: str) -> list: |
| 579 | """ |
| 580 | Fetch the artifacts. |
| 581 | |
| 582 | Returns: |
| 583 | list: The list of artifacts. |
| 584 | """ |
| 585 | artifacts = await get_artifacts() |
| 586 | return await filter_artifacts_by_os(artifacts.artifacts, os) |
| 587 | |
| 588 | |
| 589 | @threat_intel_socfortress_router.post( |
| 590 | "/ai/velociraptor-artifact-recommendation", |
| 591 | response_model=VelociraptorArtifactRecommendationResponse, |
| 592 | description="SocFortress Process Name Evaluation", |
| 593 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 594 | ) |
| 595 | async def ai_velociraptor_artifact_recommendation_socfortress( |
| 596 | request: CreateAlertRequestRoute, |
| 597 | session: AsyncSession = Depends(get_db), |
| 598 | ): |
| 599 | # Fetch alert details |
| 600 | alert_payload = await get_single_alert_details(CreateAlertRequest(index_name=request.index_name, alert_id=request.index_id)) |
| 601 | |
| 602 | assert isinstance(alert_payload, GenericAlertModel) |
| 603 | |
| 604 | os = await fetch_agent_os(request.agent_id, session) |
| 605 | |
| 606 | ai_request = VelociraptorArtifactRecommendationRequest( |
| 607 | integration="SOCFORTRESS AI", |
| 608 | alert_payload=alert_payload._source.model_dump(), |
| 609 | os=os, |
| 610 | artifacts=await fetch_artifacts(os), |
| 611 | ) |
| 612 | |
| 613 | socfortress_lookup = await socfortress_velociraptor_recommendation_lookup( |
| 614 | lincense_key=(await get_license(session)).license_key, |
| 615 | request=ai_request, |
| 616 | ) |
| 617 | |
| 618 | await create_comment( |
| 619 | CommentCreate( |
| 620 | alert_id=request.alert_id, |
| 621 | comment=f"SOCFortress AI Analysis: {socfortress_lookup.artifact_recommendations}\n\n{socfortress_lookup.general_thoughts}", |
| 622 | user_name="admin", |
| 623 | created_at=datetime.now(), |
| 624 | ), |
| 625 | db=session, |
| 626 | ) |
| 627 | return socfortress_lookup |