terminal-mcp 0.1.6

Model Context Protocol (MCP) server for long-lived shell execution.
// src/security/detect/bash/rules/destructive_disk.rs

use tree_sitter::Node;

use crate::security::detect::bash::ast::get_command_name;
use crate::security::detect::bash::utils::{
    best_hit, collect_args, command_basename, is_block_device, redirect_target, unwrap_command_where,
};
use crate::sec_bash_detector_rule_metadata;
use crate::security::detect::{EvaluateResult, Rule, Severity, ShellContext};

/// 会直接破坏块设备的命令。
const DEVICE_FORMATTERS: &[&str] = &["dd", "mkfs", "mkswap", "wipefs", "shred"];

/// 会修改分区表的命令。
const PARTITION_TOOLS: &[&str] = &["fdisk", "parted", "sfdisk"];

fn is_formatter(cmd: &str) -> bool {
    DEVICE_FORMATTERS.contains(&cmd) || cmd.starts_with("mkfs.")
}

/// 检测对块设备 / 内核内存的直接破坏性写入(dd、mkfs、shred、重定向等)。
pub struct RuleDestructiveDisk;

impl RuleDestructiveDisk {
    sec_bash_detector_rule_metadata!(
        "bash_destructive_disk",
        "Detects destructive writes to block devices or kernel memory \
         (dd to /dev/sd*, mkfs, shred, redirect to a device)",
        Severity::Low
    );
}

fn analyze_command(node: &Node, source: &[u8]) -> Option<(Severity, String)> {
    let raw_cmd = get_command_name(node, source)?;
    let cmd_name = command_basename(raw_cmd);
    let args = collect_args(node, source);
    let mut known: Vec<&str> = Vec::new();
    known.extend_from_slice(DEVICE_FORMATTERS);
    known.extend_from_slice(PARTITION_TOOLS);
    let (real_cmd, real_args) = unwrap_command_where(&cmd_name, &args, |n| {
        known.contains(&n) || n.starts_with("mkfs.")
    })?;

    if real_cmd == "dd" {
        for a in real_args {
            if let Some(of) = a.strip_prefix("of=")
                && is_block_device(of)
            {
                return Some((
                    Severity::Critical,
                    format!("command={:?} dd write to block device {:?}", raw_cmd, of),
                ));
            }
        }
        return None;
    }

    if is_formatter(real_cmd) {
        for a in real_args {
            if is_block_device(a) {
                return Some((
                    Severity::Critical,
                    format!("command={:?} {} targets block device {:?}", raw_cmd, real_cmd, a),
                ));
            }
        }
        return None;
    }

    if PARTITION_TOOLS.contains(&real_cmd) {
        for a in real_args {
            if is_block_device(a) {
                return Some((
                    Severity::High,
                    format!("command={:?} {} targets block device {:?}", raw_cmd, real_cmd, a),
                ));
            }
        }
    }

    None
}

fn analyze_redirect(node: &Node, source: &[u8]) -> Option<(Severity, String)> {
    let (op, dest) = redirect_target(node, source)?;
    if is_block_device(dest) {
        return Some((
            Severity::Critical,
            format!("{} write to block device {:?}", op, dest),
        ));
    }
    None
}

fn analyze(node: &Node, source: &[u8]) -> Option<(Severity, String)> {
    match node.kind() {
        "command" => analyze_command(node, source),
        "file_redirect" => analyze_redirect(node, source),
        _ => None,
    }
}

#[async_trait::async_trait]
impl Rule for RuleDestructiveDisk {
    fn meta(&self) -> &crate::security::detect::RuleMetadata {
        Self::get_meta()
    }

    async fn evaluate(
        &self,
        _data: &str,
        ctx: &ShellContext,
    ) -> anyhow::Result<EvaluateResult> {
        let best = best_hit(ctx, analyze).await?;
        Ok(match best {
            Some((sev, evidence)) => EvaluateResult::hit_with_severity(evidence, sev),
            None => EvaluateResult::Miss,
        })
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;

    use super::*;
    use crate::security::detect::bash::BashDetector;
    use crate::security::detect::{DetectResult, Detector, ShellContext};

    fn get_detector() -> BashDetector {
        let ctx = ShellContext::new("/bin/bash", HashMap::new(), 100);
        BashDetector::new(ctx, 4096)
    }

    async fn expect_hit(detector: &BashDetector, payload: &str, expected: Severity) {
        let res = detector.detect(payload.to_string(), false, true).await;
        match &res {
            DetectResult::ThreatDetected(hits) => {
                let hit = hits
                    .iter()
                    .find(|h| h.rule_meta.name == "bash_destructive_disk")
                    .unwrap_or_else(|| {
                        panic!("payload {:?} did not trigger bash_destructive_disk: {:#?}", payload, res)
                    });
                assert_eq!(
                    hit.final_severity, expected,
                    "payload {:?}, evidence {:?}",
                    payload, hit.evidence
                );
            }
            _ => panic!("payload {:?} expected ThreatDetected, got {:#?}", payload, res),
        }
    }

    async fn expect_safe(detector: &BashDetector, payload: &str) {
        let res = detector.detect(payload.to_string(), false, true).await;
        let hit = match &res {
            DetectResult::ThreatDetected(hits) => {
                hits.iter().find(|h| h.rule_meta.name == "bash_destructive_disk")
            }
            _ => None,
        };
        assert!(
            hit.is_none(),
            "payload {:?} should not trigger bash_destructive_disk, got {:#?}",
            payload, res
        );
    }

    #[tokio::test]
    async fn test_dd_device_critical() {
        let d = get_detector();
        expect_hit(&d, "dd if=/dev/zero of=/dev/sda bs=1M", Severity::Critical).await;
        expect_hit(&d, "dd if=/dev/random of=/dev/nvme0n1", Severity::Critical).await;
        expect_hit(&d, "dd of=/dev/sdb1", Severity::Critical).await;
        expect_hit(&d, "sudo dd if=/dev/zero of=/dev/sda", Severity::Critical).await;
        expect_hit(&d, "dd if=/dev/zero of=/dev/mem", Severity::Critical).await;
        expect_hit(&d, "timeout 5 dd if=/dev/zero of=/dev/vda", Severity::Critical).await;
    }

    #[tokio::test]
    async fn test_format_device_critical() {
        let d = get_detector();
        expect_hit(&d, "mkfs.ext4 /dev/sdb", Severity::Critical).await;
        expect_hit(&d, "mkfs -t ext4 /dev/sdc", Severity::Critical).await;
        expect_hit(&d, "mkswap /dev/sdd", Severity::Critical).await;
        expect_hit(&d, "wipefs -a /dev/sde", Severity::Critical).await;
        expect_hit(&d, "shred /dev/sdf", Severity::Critical).await;
    }

    #[tokio::test]
    async fn test_partition_high() {
        let d = get_detector();
        expect_hit(&d, "fdisk /dev/sda", Severity::High).await;
        expect_hit(&d, "parted /dev/sdb mklabel gpt", Severity::High).await;
    }

    #[tokio::test]
    async fn test_redirect_device_critical() {
        let d = get_detector();
        expect_hit(&d, "echo WIPE > /dev/sda", Severity::Critical).await;
        expect_hit(&d, "echo hi 2>> /dev/sdb", Severity::Critical).await;
    }

    #[tokio::test]
    async fn test_disk_safe() {
        let d = get_detector();
        expect_safe(&d, "dd if=/dev/zero of=/tmp/img bs=1M count=10").await;
        expect_safe(&d, "dd if=/dev/sda of=/tmp/backup.img").await;
        expect_safe(&d, "shred /tmp/secret").await;
        expect_safe(&d, "mkfs.ext4 /tmp/test.img").await;
        expect_safe(&d, "ls /dev/sda").await;
        expect_safe(&d, "echo hi > /dev/null").await;
    }
}