#!/usr/bin/env python
#
# tyr/tools/run_tyr_response.py - MCP Tool for automated threat containment
#
import os
import json
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
    }