lingshu-tools 0.10.0

//! Shared dangerous-command approval runtime for shell-based tools.
//!
//! WHY shared module: `terminal` and `run_process` both need the same command
//! scan, approval request, session cache, and persistence behavior. Keeping
//! the flow here avoids divergent security semantics across shell tools.

use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::{Mutex, OnceLock, RwLock};

use lingshu_security::approval::{ApprovalMode, ApprovalPolicy};
use lingshu_types::ToolError;
use serde_json::json;

use crate::registry::{ApprovalRequest, ApprovalResponse, ToolContext};
use crate::smart_approval::{SmartVerdict, assess_smart_approval};

fn approval_policy() -> &'static ApprovalPolicy {
    static POLICY: OnceLock<ApprovalPolicy> = OnceLock::new();
    POLICY.get_or_init(|| ApprovalPolicy::new(ApprovalMode::Manual, Vec::new()))
}

fn allowlist_file_lock() -> &'static Mutex<()> {
    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| Mutex::new(()))
}

fn yolo_sessions() -> &'static RwLock<HashSet<String>> {
    static SESSIONS: OnceLock<RwLock<HashSet<String>>> = OnceLock::new();
    SESSIONS.get_or_init(|| RwLock::new(HashSet::new()))
}

fn allowlist_path(lingshu_home: &Path) -> std::path::PathBuf {
    lingshu_home.join("command_allowlist.json")
}

fn load_persistent_allowlist(lingshu_home: &Path) -> Vec<String> {
    let path = allowlist_path(lingshu_home);
    let Ok(raw) = std::fs::read_to_string(path) else {
        return Vec::new();
    };
    serde_json::from_str::<Vec<String>>(&raw).unwrap_or_default()
}

fn persist_allowlist_command(lingshu_home: &Path, command: &str) -> Result<(), ToolError> {
    let _guard = allowlist_file_lock()
        .lock()
        .map_err(|_| ToolError::Other("approval allowlist lock poisoned".into()))?;

    let path = allowlist_path(lingshu_home);
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| ToolError::Other(format!("failed to create allowlist directory: {e}")))?;
    }

    let mut entries: HashSet<String> = load_persistent_allowlist(lingshu_home)
        .into_iter()
        .collect();
    entries.insert(command.to_string());

    let mut sorted: Vec<String> = entries.into_iter().collect();
    sorted.sort();

    let payload = serde_json::to_string_pretty(&sorted)
        .map_err(|e| ToolError::Other(format!("failed to serialize allowlist: {e}")))?;
    std::fs::write(path, payload)
        .map_err(|e| ToolError::Other(format!("failed to persist allowlist: {e}")))?;
    Ok(())
}

fn command_preview(command: &str) -> String {
    let single_line = command.split_whitespace().collect::<Vec<_>>().join(" ");
    crate::safe_truncate(&single_line, 80).to_string()
}

fn session_key(ctx: &ToolContext) -> String {
    ctx.session_key
        .clone()
        .unwrap_or_else(|| ctx.session_id.clone())
}

pub fn yolo_enabled_for_session(session_key: &str) -> bool {
    yolo_sessions()
        .read()
        .map(|sessions| sessions.contains(session_key))
        .unwrap_or(false)
}

pub fn set_yolo_for_session(session_key: impl Into<String>, enabled: bool) {
    let session_key = session_key.into();
    if let Ok(mut sessions) = yolo_sessions().write() {
        if enabled {
            sessions.insert(session_key);
        } else {
            sessions.remove(&session_key);
        }
    }
}

// ─── computer_use per-action session approval (Hermes `_always_allow`) ───

fn computer_use_action_cache() -> &'static RwLock<HashMap<String, HashSet<String>>> {
    static CACHE: OnceLock<RwLock<HashMap<String, HashSet<String>>>> = OnceLock::new();
    CACHE.get_or_init(|| RwLock::new(HashMap::new()))
}

/// True when this destructive `computer_use` action was approved for the session.
pub fn computer_use_action_approved(session_key: &str, action: &str) -> bool {
    computer_use_action_cache()
        .read()
        .ok()
        .and_then(|cache| cache.get(session_key).map(|set| set.contains(action)))
        .unwrap_or(false)
}

/// Remember approval for one action type until session ends (Hermes `approve_session`).
pub fn approve_computer_use_action_for_session(session_key: impl Into<String>, action: &str) {
    let session_key = session_key.into();
    if let Ok(mut cache) = computer_use_action_cache().write() {
        cache
            .entry(session_key)
            .or_default()
            .insert(action.to_string());
    }
}

pub(crate) fn command_approval_reasons(ctx: &ToolContext, command: &str) -> Option<Vec<String>> {
    if yolo_enabled_for_session(&session_key(ctx)) {
        return None;
    }

    let mode = ctx.config.approval_mode;
    if mode == ApprovalMode::Off {
        return None;
    }

    let policy = approval_policy();
    policy.load_permanent_allowlist(&load_persistent_allowlist(&ctx.config.lingshu_home));

    let session_id = session_key(ctx);
    let check = policy.check_with_mode(
        mode,
        "terminal",
        &json!({ "command": command }),
        &session_id,
    );
    if check.needs_approval {
        Some(check.reasons)
    } else {
        None
    }
}

pub(crate) async fn request_command_approval(
    ctx: &ToolContext,
    command: &str,
    reasons: Vec<String>,
) -> Result<(), ToolError> {
    if yolo_enabled_for_session(&session_key(ctx)) {
        return Ok(());
    }

    let mode = ctx.config.approval_mode;
    if mode == ApprovalMode::Off {
        return Ok(());
    }

    let policy = approval_policy();
    policy.load_permanent_allowlist(&load_persistent_allowlist(&ctx.config.lingshu_home));

    let session_id = session_key(ctx);
    let check = policy.check_with_mode(
        mode,
        "terminal",
        &json!({ "command": command }),
        &session_id,
    );
    if !check.needs_approval {
        return Ok(());
    }

    if mode == ApprovalMode::Smart {
        let description = reasons.join("; ");
        match assess_smart_approval(command, &description, ctx).await {
            SmartVerdict::Approve => return Ok(()),
            SmartVerdict::Deny => {
                return Err(ToolError::PermissionDenied(
                    "Command denied by smart approval policy.".into(),
                ));
            }
            SmartVerdict::Escalate => {}
        }
    }

    let Some(tx) = &ctx.approval_tx else {
        return Err(ToolError::PermissionDenied(format!(
            "Dangerous command requires approval: {}",
            reasons.join(", ")
        )));
    };

    let (resp_tx, resp_rx) = tokio::sync::oneshot::channel::<ApprovalResponse>();
    tx.send(ApprovalRequest {
        command: command_preview(command),
        full_command: command.to_string(),
        reasons,
        response_tx: resp_tx,
    })
    .map_err(|_| {
        ToolError::PermissionDenied(
            "Dangerous command requires approval, but no interactive approver is available.".into(),
        )
    })?;

    let response = tokio::select! {
        _ = ctx.cancel.cancelled() => {
            return Err(ToolError::Other("Interrupted by user".into()));
        }
        result = resp_rx => result.map_err(|_| ToolError::PermissionDenied(
            "Approval request was cancelled before a decision was received.".into(),
        ))?
    };

    match response {
        ApprovalResponse::Once => Ok(()),
        ApprovalResponse::Session => {
            policy.approve_for_session(command, &session_id);
            Ok(())
        }
        ApprovalResponse::Always => {
            policy.approve_for_session(command, &session_id);
            policy.approve_permanently(command);
            persist_allowlist_command(&ctx.config.lingshu_home, command)?;
            Ok(())
        }
        ApprovalResponse::Deny => Err(ToolError::PermissionDenied(
            "Command denied by user approval policy.".into(),
        )),
    }
}

#[cfg(test)]
mod computer_use_approval_tests {
    use super::*;

    #[test]
    fn session_approval_caches_per_action() {
        let session = "test-session-computer-use";
        assert!(!computer_use_action_approved(session, "key"));
        approve_computer_use_action_for_session(session, "key");
        assert!(computer_use_action_approved(session, "key"));
        assert!(!computer_use_action_approved(session, "click"));
    }
}