|
30 | 30 | api_config = Config("api") |
31 | 31 |
|
32 | 32 | # Configuration from Environment or Config File |
33 | | -# Run with: CAPE_API_URL=http://127.0.0.1:8000/apiv2 CAPE_API_TOKEN=your_token python3 web/mcp_server.py |
| 33 | +# Run with: CAPE_API_URL=http://127.0.0.1:8000/apiv2 CAPE_API_TOKEN=your_token poetry run python mcp/server.py |
34 | 34 | API_URL = os.environ.get("CAPE_API_URL") |
35 | 35 | if not API_URL: |
36 | 36 | # Try to get from api.conf [api] url |
@@ -344,26 +344,69 @@ async def submit_static( |
344 | 344 |
|
345 | 345 | # --- Task Management & Search --- |
346 | 346 |
|
| 347 | +def get_lean_cape_report(raw_cape_json): |
| 348 | + """Filters a 50MB CAPE report down to a 500-token LLM payload.""" |
| 349 | + return { |
| 350 | + "score": raw_cape_json.get("info", {}).get("score", 0), |
| 351 | + "family": raw_cape_json.get("malfamily") or raw_cape_json.get("detections", {}).get("family") or "Unknown", |
| 352 | + "extracted_configs": raw_cape_json.get("CAPE", []), |
| 353 | + "high_severity_signatures": [ |
| 354 | + {"name": sig["name"], "desc": sig["description"]} |
| 355 | + for sig in raw_cape_json.get("signatures", []) |
| 356 | + if isinstance(sig, dict) and sig.get("severity", 0) >= 3 |
| 357 | + ], |
| 358 | + "network": { |
| 359 | + "domains": [d["domain"] for d in raw_cape_json.get("network", {}).get("domains", [])] if isinstance(raw_cape_json.get("network", {}).get("domains"), list) else [], |
| 360 | + "http_uris": [h["uri"] for h in raw_cape_json.get("network", {}).get("http", [])] if isinstance(raw_cape_json.get("network", {}).get("http"), list) else [], |
| 361 | + }, |
| 362 | + "indicators": { |
| 363 | + "mutexes": raw_cape_json.get("behavior", {}).get("summary", {}).get("mutexes", []) if isinstance(raw_cape_json.get("behavior", {}).get("summary"), dict) else [], |
| 364 | + "commands": raw_cape_json.get("behavior", {}).get("summary", {}).get("executed_commands", []) if isinstance(raw_cape_json.get("behavior", {}).get("summary"), dict) else [] |
| 365 | + } |
| 366 | + } |
| 367 | + |
| 368 | +def _apply_lean_report(result): |
| 369 | + if isinstance(result, dict): |
| 370 | + if result.get("error") is False and "data" in result: |
| 371 | + if isinstance(result["data"], list): |
| 372 | + result["data"] = [get_lean_cape_report(item) for item in result["data"]] |
| 373 | + elif isinstance(result["data"], dict): |
| 374 | + result["data"] = get_lean_cape_report(result["data"]) |
| 375 | + elif "info" in result: |
| 376 | + return get_lean_cape_report(result) |
| 377 | + elif isinstance(result, list): |
| 378 | + return [get_lean_cape_report(item) for item in result] |
| 379 | + return result |
| 380 | + |
347 | 381 | @mcp_tool("tasksearch") |
348 | | -async def search_task(hash_value: str, token: str = "") -> str: |
| 382 | +async def search_task(hash_value: str, lean: bool = True, token: str = "") -> str: |
349 | 383 | """Search for tasks by MD5, SHA1, or SHA256.""" |
| 384 | + if not re.match(r"^[a-fA-F0-9]+$", hash_value): |
| 385 | + return json.dumps({"error": True, "message": "Invalid hash value provided. Only hexadecimal characters are allowed."}, indent=2) |
| 386 | + |
350 | 387 | algo = "md5" |
351 | 388 | if len(hash_value) == 40: |
352 | 389 | algo = "sha1" |
353 | 390 | elif len(hash_value) == 64: |
354 | 391 | algo = "sha256" |
355 | 392 |
|
356 | 393 | result = await _request("GET", f"tasks/search/{algo}/{hash_value}/", token=token) |
| 394 | + if lean: |
| 395 | + result = _apply_lean_report(result) |
357 | 396 | return json.dumps(result, indent=2) |
358 | 397 |
|
359 | 398 | @mcp_tool("extendedtasksearch") |
360 | | -async def extended_search(option: str, argument: str, token: str = "") -> str: |
| 399 | +async def extended_search(option: str, argument: str, lean: bool = True, token: str = "") -> str: |
361 | 400 | """ |
362 | 401 | Search tasks using extended options. |
363 | 402 | Options include: id, name, type, string, ssdeep, crc32, file, command, resolvedapi, key, mutex, domain, ip, signature, signame, etc. |
364 | 403 | """ |
365 | 404 | data = {"option": option, "argument": argument} |
| 405 | + if lean: |
| 406 | + data["lean"] = True |
366 | 407 | result = await _request("POST", "tasks/extendedsearch/", token=token, data=data) |
| 408 | + if lean: |
| 409 | + result = _apply_lean_report(result) |
367 | 410 | return json.dumps(result, indent=2) |
368 | 411 |
|
369 | 412 | @mcp_tool("extendedtasksearch") |
@@ -430,7 +473,25 @@ async def get_statistics(days: int = 7, token: str = "") -> str: |
430 | 473 |
|
431 | 474 | @mcp_tool("taskreport") |
432 | 475 | async def get_task_report(task_id: int, format: str = "json", token: str = "") -> str: |
433 | | - """Get the analysis report for a task (json, lite, maec, metadata).""" |
| 476 | + """Get the analysis report for a task (json, lite, maec, metadata, lean).""" |
| 477 | + allowed_formats = {"json", "lite", "maec", "metadata", "lean"} |
| 478 | + if format not in allowed_formats: |
| 479 | + return json.dumps({"error": True, "message": f"Invalid format provided. Allowed formats: {', '.join(allowed_formats)}"}, indent=2) |
| 480 | + |
| 481 | + if format == "lean": |
| 482 | + data = {"option": "id", "argument": str(task_id), "lean": True} |
| 483 | + result = await _request("POST", "tasks/extendedsearch/", token=token, data=data) |
| 484 | + |
| 485 | + # Extract the single task report from the search results |
| 486 | + if isinstance(result, dict) and not result.get("error") and isinstance(result.get("data"), list): |
| 487 | + if len(result["data"]) > 0: |
| 488 | + result["data"] = result["data"][0] |
| 489 | + else: |
| 490 | + result = {"error": True, "message": "Task report not found via lean search."} |
| 491 | + |
| 492 | + result = _apply_lean_report(result) |
| 493 | + return json.dumps(result, indent=2) |
| 494 | + |
434 | 495 | result = await _request("GET", f"tasks/get/report/{task_id}/{format}/", token=token) |
435 | 496 | return json.dumps(result, indent=2) |
436 | 497 |
|
@@ -516,11 +577,15 @@ async def download_task_fullmemory(task_id: int, destination: str, token: str = |
516 | 577 | @mcp_tool("fileview") |
517 | 578 | async def view_file(hash_value: str, hash_type: str = "sha256", token: str = "") -> str: |
518 | 579 | """View information about a file in the database.""" |
| 580 | + if not re.match(r"^[a-fA-F0-9]+$", hash_value): |
| 581 | + return json.dumps({"error": True, "message": "Invalid hash value provided. Only hexadecimal characters are allowed."}, indent=2) |
519 | 582 | return await _request("GET", f"files/view/{hash_type}/{hash_value}/", token=token) |
520 | 583 |
|
521 | 584 | @mcp_tool("sampledl") |
522 | 585 | async def download_sample(hash_value: str, destination: str, hash_type: str = "sha256", token: str = "") -> str: |
523 | 586 | """Download a sample from the database.""" |
| 587 | + if not re.match(r"^[a-fA-F0-9]+$", hash_value): |
| 588 | + return json.dumps({"error": True, "message": "Invalid hash value provided. Only hexadecimal characters are allowed."}, indent=2) |
524 | 589 | return await _download_file(f"files/get/{hash_type}/{hash_value}/", destination, f"{hash_value}.bin", token=token) |
525 | 590 |
|
526 | 591 | @mcp_tool("machinelist") |
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?v=4&size=40){kind=avatar}
0 commit comments