90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
#!/usr/bin/env python
|
|
#
|
|
# tyr/tools/run_tyr_response.py - MCP Tool for automated threat containment
|
|
#
|
|
import os
|
|
import subprocess
|
|
|
|
SPIRE_SERVER_PATH = "./spire-1.15.3/bin/spire-server"
|
|
|
|
|
|
def _run_command(cmd_list, description):
|
|
"""Helper to run a command and return its status and output."""
|
|
try:
|
|
result = subprocess.run(
|
|
cmd_list, capture_output=True, text=True, check=True, timeout=30)
|
|
return {"status": "SUCCESS", "description": description, "output": result.stdout or result.stderr}
|
|
except subprocess.CalledProcessError as e:
|
|
return {"status": "FAIL", "description": description, "error": e.stderr or e.stdout}
|
|
except Exception as e:
|
|
return {"status": "ERROR", "description": description, "message": str(e)}
|
|
|
|
|
|
def _revoke_spiffe_id(target):
|
|
spiffe_id = target.get("spiffe_id")
|
|
if not spiffe_id:
|
|
raise ValueError("Missing 'spiffe_id' for revoke_spiffe_id action.")
|
|
cmd = [SPIRE_SERVER_PATH, "entry", "delete", "-spiffeID", spiffe_id]
|
|
return _run_command(cmd, f"Attempted to delete SPIFFE ID: {spiffe_id}")
|
|
|
|
|
|
def _disable_secret_version(target):
|
|
secret = target.get("secret")
|
|
version = target.get("version", "latest")
|
|
if not secret:
|
|
raise ValueError("Missing 'secret' for disable_secret_version action.")
|
|
cmd = ["gcloud", "secrets", "versions", "disable",
|
|
version, "--secret", secret, "--quiet"]
|
|
return _run_command(cmd, f"Attempted to disable version '{version}' of secret '{secret}'.")
|
|
|
|
|
|
def _revoke_iam_member(target):
|
|
project_id = os.environ.get(
|
|
"GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
|
|
member = target.get("member")
|
|
role = target.get("role")
|
|
if not (member and role):
|
|
raise ValueError(
|
|
"Missing 'member' and/or 'role' for revoke_iam_member action.")
|
|
cmd = ["gcloud", "projects", "remove-iam-policy-binding",
|
|
project_id, "--member", member, "--role", role, "--quiet"]
|
|
return _run_command(cmd, f"Attempted to remove role '{role}' from member '{member}'.")
|
|
|
|
|
|
def _quarantine_ip(target):
|
|
ip_address = target.get("ip_address")
|
|
if not ip_address:
|
|
raise ValueError("Missing 'ip_address' for quarantine_ip action.")
|
|
# This action would add the IP to a deny rule in a real implementation.
|
|
return {"status": "NOT_IMPLEMENTED", "message": f"Quarantine action for {ip_address} is a placeholder and was not executed."}
|
|
|
|
|
|
ACTION_MAP = {
|
|
"revoke_spiffe_id": _revoke_spiffe_id,
|
|
"disable_secret_version": _disable_secret_version,
|
|
"revoke_iam_member": _revoke_iam_member,
|
|
"quarantine_ip": _quarantine_ip,
|
|
}
|
|
|
|
|
|
def run_tyr_response(p: dict) -> dict:
|
|
"""Executes a predefined automated threat containment action."""
|
|
action = p.get("action")
|
|
target = p.get("target")
|
|
|
|
if not action or not target:
|
|
raise ValueError("Missing required parameters: 'action' and 'target'")
|
|
|
|
handler = ACTION_MAP.get(action)
|
|
if not handler:
|
|
raise ValueError(
|
|
f"Invalid action specified: {action}. Must be one of {list(ACTION_MAP.keys())}")
|
|
|
|
result = handler(target)
|
|
|
|
return {
|
|
"action_taken": action,
|
|
"target_resource": target,
|
|
"execution_result": result
|
|
}
|