zesdex/src/app/harness.rs

//! Tool-call gating: decides whether a risky tool call is allowed to run
//! before it executes.

/// Outcome of gating a tool call: whether it's allowed to run.
#[derive(Debug, Clone, PartialEq)]
pub enum Verdict {
    Allow,
    Block(String),
}

/// Gatekeeper that decides whether a tool call may proceed before execution.
pub struct Harness;

impl Harness {
    /// Decide whether a tool call is allowed to execute.
    ///
    /// Flow: if the tool isn't flagged risky, allow immediately → basic
    /// content checks (path traversal in paths AND command args) →
    /// workspace-root validation for output paths → classify.
    ///
    /// Return: `Verdict::Allow` or `Verdict::Block(reason)`.
    pub fn gate_tool_call(
        tool_name: &str,
        args: &serde_json::Value,
        workspace_roots: &[&std::path::Path],
    ) -> Verdict {

        if !crate::tool::tool_is_risky(tool_name) {
            return Verdict::Allow;
        }
        // Path traversal check for file-mutating tools.
        if matches!(tool_name, "write" | "edit" | "delete") {
            if let Some(path) = args.get("path").and_then(|v| v.as_str()) {
                if path.contains("..") {
                    return Verdict::Block("path traversal detected in 'path' argument".to_string());
                }
            }
        }
        // Path traversal and dangerous content check for bash commands.
        if tool_name == "bash" {
            let cmd = args.get("command").and_then(|v| v.as_str()).unwrap_or("");
            if cmd.contains("..") {
                return Verdict::Block("path traversal detected in bash command".to_string());
            }
            let dangerous_patterns = [
                "rm -rf /", "rm -rf --no-preserve-root",
                "rm -rf ~", "rm -fr /", "mkfs.", "dd if=",
                ":(){", "> /dev/sda", "chmod -R 000 /",
            ];
            for pat in &dangerous_patterns {
                if cmd.contains(pat) {
                    return Verdict::Block(format!("destructive command pattern blocked: {}", pat));
                }
            }
        }
        if let Some(out_path) = Self::find_output_path(tool_name, args) {
            if !workspace_roots.is_empty()
                && !out_path.starts_with("/tmp")
                && !out_path.is_absolute()
            {
                let allowed = workspace_roots.iter().any(|r| out_path.starts_with(r));
                if !allowed {
                    return Verdict::Block(format!(
                        "output path '{:?}' is outside all workspace roots", out_path
                    ));
                }
            }
        }
        Self::classify(tool_name)
    }

    /// Extract a candidate output path from a tool call, if one exists.
    ///
    /// Used to verify that writes and file mutations stay inside workspace roots.
    fn find_output_path(tool_name: &str, args: &serde_json::Value) -> Option<std::path::PathBuf> {
        match tool_name {
            "write" | "edit" | "delete" | "read" => {
                args.get("path").and_then(|v| v.as_str()).map(std::path::PathBuf::from)
            }
            "bash" => {
                let cmd = args.get("command").and_then(|v| v.as_str())?;
                let lower = cmd.to_lowercase();
                for prefix in &["cp ", "mv ", "install ", "ln -s ", "cat >", "cat >>"] {
                    if let Some(rest) = lower.strip_prefix(prefix) {
                        if let Some(target) = rest.split_whitespace().last() {
                            if !target.starts_with('-') {
                                return Some(std::path::PathBuf::from(target));
                            }
                        }
                    }
                }
                None
            }
            _ => None,
        }
    }

    fn classify(cmd: &str) -> Verdict {
        match cmd {
            "bash" | "write" | "edit" | "delete" | "git_operator" => Verdict::Allow,
            _ => Verdict::Allow,
        }
    }

}

impl Default for Harness {
    fn default() -> Self {
        Harness
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn parse_verdict(text: &str) -> Option<Verdict> {
        let trimmed = text.trim();
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(trimmed) {
            if let Some(verdict) = v.get("verdict").and_then(|v| v.as_str()) {
                return match verdict.to_lowercase().as_str() {
                    "allow" => Some(Verdict::Allow),
                    "block" => Some(Verdict::Block(
                        v.get("reason").and_then(|r| r.as_str()).unwrap_or("blocked").to_string()
                    )),
                    _ => None,
                };
            }
        }
        for line in trimmed.lines() {
            let l = line.trim().to_lowercase();
            if l.starts_with("verdict: allow") {
                return Some(Verdict::Allow);
            }
            if l.starts_with("verdict: block") {
                let reason = line.split_once(':').map(|x| x.1).unwrap_or("blocked").trim().to_string();
                return Some(Verdict::Block(reason));
            }
        }
        if trimmed.to_lowercase().contains("allow") {
            return Some(Verdict::Allow);
        }
        if trimmed.to_lowercase().contains("block") {
            return Some(Verdict::Block("blocked by classifier".to_string()));
        }
        None
    }

    #[test]
    fn test_classify_always_allows() {
        assert_eq!(Harness::classify("write"), Verdict::Allow);
    }

    #[test]
    fn test_gate_tool_non_risky_always_allows() {
        let roots: &[&std::path::Path] = &[];
        let result = Harness::gate_tool_call("read", &json!({"path": "test.txt"}), roots);
        assert_eq!(result, Verdict::Allow);
    }


    #[test]
    fn test_parse_verdict_json_allow() {
        let v = parse_verdict(r#"{"verdict": "allow"}"#);
        assert_eq!(v, Some(Verdict::Allow));
    }

    #[test]
    fn test_parse_verdict_json_block() {
        let v = parse_verdict(r#"{"verdict": "block", "reason": "dangerous operation"}"#);
        assert_eq!(v, Some(Verdict::Block("dangerous operation".to_string())));
    }

    #[test]
    fn test_parse_verdict_text_allow() {
        let v = parse_verdict("Verdict: Allow");
        assert_eq!(v, Some(Verdict::Allow));
    }

    #[test]
    fn test_parse_verdict_text_block() {
        let v = parse_verdict("Verdict: Block - this operation is not allowed");
        assert!(matches!(v, Some(Verdict::Block(_))));
    }

    #[test]
    fn test_parse_verdict_fallback_allow() {
        let v = parse_verdict("I think we should allow this operation");
        assert_eq!(v, Some(Verdict::Allow));
    }

    #[test]
    fn test_parse_verdict_fallback_block() {
        let v = parse_verdict("This request should be blocked");
        assert!(matches!(v, Some(Verdict::Block(_))));
    }

    #[test]
    fn test_parse_verdict_unparseable() {
        let v = parse_verdict("completely unrelated text with no keywords");
        assert_eq!(v, None);
    }
}