securegit 0.8.5

use rmcp::model::CallToolRequestParams;
use rmcp::transport::child_process::TokioChildProcess;
use rmcp::ServiceExt;
use tokio::process::Command;

const DEFAULT_BINARY: &str = "armyknife-llm-redteam-mcp";

/// MCP-to-MCP bridge to armyknife-llm-redteam.
///
/// Spawns a fresh connection for each tool call.
/// Gracefully degrades when the binary is not installed.
pub struct RedteamBridge {
    binary: String,
}

/// Distinguishes "binary not found" (skip) from "call failed" (error).
#[derive(Debug)]
pub enum BridgeError {
    /// Binary not installed — caller should skip LLM security checks.
    NotInstalled(String),
    /// Binary found but a tool call failed — surface to user.
    CallFailed(String),
}

impl std::fmt::Display for BridgeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotInstalled(msg) => write!(f, "redteam not installed: {}", msg),
            Self::CallFailed(msg) => write!(f, "redteam call failed: {}", msg),
        }
    }
}

impl std::error::Error for BridgeError {}

impl RedteamBridge {
    /// Create a new bridge.
    pub fn new(binary: Option<String>) -> Self {
        let binary = std::env::var("SECUREGIT_REDTEAM_BIN")
            .ok()
            .or(binary)
            .unwrap_or_else(|| DEFAULT_BINARY.to_string());

        Self { binary }
    }

    /// Check if the binary exists on PATH without spawning it.
    pub fn is_available(&self) -> bool {
        which::which(&self.binary).is_ok()
    }

    /// Call a tool on the redteam MCP server. Returns the JSON result.
    async fn call_tool(
        &self,
        tool_name: &str,
        params: serde_json::Value,
    ) -> Result<String, BridgeError> {
        if !self.is_available() {
            return Err(BridgeError::NotInstalled(format!(
                "{} not found on PATH",
                self.binary
            )));
        }

        // Spawn a fresh connection for this call
        let child = TokioChildProcess::new(Command::new(&self.binary))
            .map_err(|e| BridgeError::CallFailed(format!("Failed to spawn: {}", e)))?;

        let service = ()
            .serve(child)
            .await
            .map_err(|e| BridgeError::CallFailed(format!("Failed to initialize: {}", e)))?;

        let request = CallToolRequestParams {
            meta: None,
            name: tool_name.to_string().into(),
            arguments: params.as_object().cloned(),
            task: None,
        };

        let result = service
            .call_tool(request)
            .await
            .map_err(|e| BridgeError::CallFailed(format!("{}: {}", tool_name, e)))?;

        // Shutdown the service
        service.cancel().await.ok();

        // Convert result to JSON string
        let text = serde_json::to_string(&result).unwrap_or_else(|_| format!("{:?}", result));

        Ok(text)
    }

    // ---- Public tool wrappers ----

    /// Scan an MCP server for tool poisoning, shadowing, sandbox escape.
    pub async fn scan_mcp_server(
        &self,
        command: &str,
        args: &[String],
    ) -> Result<String, BridgeError> {
        self.call_tool(
            "redteam_scan_mcp",
            serde_json::json!({
                "command": command,
                "args": args,
            }),
        )
        .await
    }

    /// Check a prompt string for injection risks.
    pub async fn firewall_check(&self, input: &str) -> Result<String, BridgeError> {
        self.call_tool(
            "redteam_firewall_check",
            serde_json::json!({ "input": input }),
        )
        .await
    }

    /// Pin MCP tool definitions at known-good state.
    pub async fn pin_tools(
        &self,
        command: &str,
        args: &[String],
        output_path: Option<&str>,
    ) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({
            "command": command,
            "args": args,
        });
        if let Some(path) = output_path {
            params["output_path"] = serde_json::json!(path);
        }
        self.call_tool("redteam_pin", params).await
    }

    /// Verify pinned MCP tool definitions (rug-pull detection).
    pub async fn verify_pins(
        &self,
        command: &str,
        args: &[String],
        pins_path: &str,
    ) -> Result<String, BridgeError> {
        self.call_tool(
            "redteam_verify",
            serde_json::json!({
                "command": command,
                "args": args,
                "pins_path": pins_path,
            }),
        )
        .await
    }

    /// Full LLM endpoint scan (slow — network calls).
    pub async fn scan_endpoint(
        &self,
        endpoint: &str,
        model: &str,
        provider: Option<&str>,
        format: Option<&str>,
    ) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({
            "endpoint": endpoint,
            "model": model,
        });
        if let Some(p) = provider {
            params["provider"] = serde_json::json!(p);
        }
        if let Some(f) = format {
            params["format"] = serde_json::json!(f);
        }
        self.call_tool("redteam_scan", params).await
    }

    /// Pipeline gate scan (slow — network calls).
    pub async fn pipeline_scan(
        &self,
        model_spec: &str,
        output_dir: Option<&str>,
    ) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({ "model": model_spec });
        if let Some(dir) = output_dir {
            params["output"] = serde_json::json!(dir);
        }
        self.call_tool("redteam_pipeline_scan", params).await
    }

    /// Query cached findings from last scan.
    pub async fn findings(&self, min_severity: Option<&str>) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({});
        if let Some(sev) = min_severity {
            params["min_severity"] = serde_json::json!(sev);
        }
        self.call_tool("redteam_findings", params).await
    }

    /// Pipeline harden: generate training data + fine-tune (supports HF cloud).
    #[allow(clippy::too_many_arguments)]
    pub async fn pipeline_harden(
        &self,
        model: &str,
        findings: &str,
        output: Option<&str>,
        format: Option<&str>,
        refusal_style: Option<&str>,
        epochs: Option<u32>,
        train_on: Option<&str>,
        hf_org: Option<&str>,
        hf_hardware: Option<&str>,
    ) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({
            "model": model,
            "findings": findings,
        });
        if let Some(o) = output {
            params["output"] = serde_json::json!(o);
        }
        if let Some(f) = format {
            params["format"] = serde_json::json!(f);
        }
        if let Some(r) = refusal_style {
            params["refusal_style"] = serde_json::json!(r);
        }
        if let Some(e) = epochs {
            params["epochs"] = serde_json::json!(e);
        }
        if let Some(t) = train_on {
            params["train_on"] = serde_json::json!(t);
        }
        if let Some(o) = hf_org {
            params["hf_org"] = serde_json::json!(o);
        }
        if let Some(h) = hf_hardware {
            params["hf_hardware"] = serde_json::json!(h);
        }
        self.call_tool("redteam_pipeline_harden", params).await
    }

    /// Pipeline verify: targeted re-scan of hardened model against original findings.
    pub async fn pipeline_verify(
        &self,
        original_findings: &str,
        hardened_model: &str,
        original_model: Option<&str>,
        output: Option<&str>,
        min_fix_rate: Option<f64>,
        fail_on_regression: Option<bool>,
    ) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({
            "original_findings": original_findings,
            "hardened_model": hardened_model,
        });
        if let Some(m) = original_model {
            params["original_model"] = serde_json::json!(m);
        }
        if let Some(o) = output {
            params["output"] = serde_json::json!(o);
        }
        if let Some(r) = min_fix_rate {
            params["min_fix_rate"] = serde_json::json!(r);
        }
        if let Some(f) = fail_on_regression {
            params["fail_on_regression"] = serde_json::json!(f);
        }
        self.call_tool("redteam_pipeline_verify", params).await
    }

    /// Pipeline publish: publish hardened model to HuggingFace.
    pub async fn pipeline_publish(
        &self,
        model_path: &str,
        report_path: &str,
        model_name: &str,
        hf_org: Option<&str>,
    ) -> Result<String, BridgeError> {
        let mut params = serde_json::json!({
            "model_path": model_path,
            "report_path": report_path,
            "model_name": model_name,
        });
        if let Some(o) = hf_org {
            params["hf_org"] = serde_json::json!(o);
        }
        self.call_tool("redteam_pipeline_publish", params).await
    }
}