punch_types/

taint.rs

//! Taint tracking and shell bleed detection.
//!
//! Every piece of data entering the ring carries taint labels — markers that
//! track where the data came from and how sensitive it is. When a fighter
//! attempts a shell move, the **bleed detector** scans the command for leaked
//! secrets before the punch lands. If a tainted move bleeds confidential data,
//! it gets blocked before it can do damage.

use chrono::{DateTime, Utc};
use regex::Regex;
use serde::{Deserialize, Serialize};

// ---------------------------------------------------------------------------
// Taint source and sensitivity
// ---------------------------------------------------------------------------

/// Where a piece of data originated — every tainted move has an origin story.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum TaintSource {
    /// Data entered directly by a human operator.
    UserInput,
    /// Output from a tool (move) execution. The string identifies the tool.
    ToolOutput(String),
    /// Generated by an LLM during the bout.
    LlmGenerated,
    /// Received from an external API. The string identifies the API.
    ExternalApi(String),
    /// Read from the filesystem. The string is the file path.
    FileSystem(String),
    /// Sourced from the process environment.
    Environment,
}

/// Sensitivity level for tainted data — determines how hard the bleed
/// detection hits when it finds a leak.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum Sensitivity {
    /// Public data — no bleed risk.
    Public,
    /// Internal data — log a warning but let the move land.
    Internal,
    /// Confidential data — block the move.
    Confidential,
    /// Secret data — block the move and raise an alert.
    Secret,
}

impl std::fmt::Display for Sensitivity {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Public => write!(f, "public"),
            Self::Internal => write!(f, "internal"),
            Self::Confidential => write!(f, "confidential"),
            Self::Secret => write!(f, "secret"),
        }
    }
}

// ---------------------------------------------------------------------------
// Taint labels and tainted values
// ---------------------------------------------------------------------------

/// A taint label attached to a value — records the origin, sensitivity, and
/// how the taint propagated through the ring.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaintLabel {
    /// Where this data came from.
    pub source: TaintSource,
    /// How sensitive this data is.
    pub sensitivity: Sensitivity,
    /// When the taint was first applied.
    pub timestamp: DateTime<Utc>,
    /// Chain of operations that propagated this taint (e.g. tool names,
    /// transformation steps).
    pub propagation_chain: Vec<String>,
}

/// A value carrying taint labels — knows where it has been and what it
/// carries from the ring.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaintedValue {
    /// The raw string value.
    pub value: String,
    /// All taint labels attached to this value.
    pub labels: Vec<TaintLabel>,
}

// ---------------------------------------------------------------------------
// TaintTracker
// ---------------------------------------------------------------------------

/// Tracks taint propagation across tool calls in a bout.
///
/// The tracker keeps a registry of tainted values and their labels. When data
/// flows from one move to another, `propagate` extends the chain so we can
/// trace the full path a tainted value took through the ring.
#[derive(Debug, Clone, Default)]
pub struct TaintTracker {
    /// All registered tainted values.
    values: Vec<TaintedValue>,
}

impl TaintTracker {
    /// Create a fresh tracker — no tainted moves recorded yet.
    pub fn new() -> Self {
        Self { values: Vec::new() }
    }

    /// Register a value as tainted with the given source and sensitivity.
    ///
    /// This marks the value as carrying sensitive data from the specified
    /// origin. Like marking a fighter's gloves before a bout.
    pub fn taint(&mut self, value: &str, source: TaintSource, sensitivity: Sensitivity) {
        // Check if this exact value is already tracked.
        for tv in &mut self.values {
            if tv.value == value {
                tv.labels.push(TaintLabel {
                    source,
                    sensitivity,
                    timestamp: Utc::now(),
                    propagation_chain: Vec::new(),
                });
                return;
            }
        }

        self.values.push(TaintedValue {
            value: value.to_string(),
            labels: vec![TaintLabel {
                source,
                sensitivity,
                timestamp: Utc::now(),
                propagation_chain: Vec::new(),
            }],
        });
    }

    /// Check if a value (or any tainted substring) is tainted.
    ///
    /// Returns all matching taint labels. Uses simple substring matching —
    /// if a tainted value appears anywhere inside the checked string, or if
    /// the checked string appears inside a tainted value, the taint is
    /// detected. No fancy footwork needed.
    pub fn check_taint(&self, value: &str) -> Vec<&TaintLabel> {
        let mut labels = Vec::new();
        for tv in &self.values {
            // Match if the tainted value is a substring of the checked value,
            // or the checked value is a substring of the tainted value.
            if value.contains(&tv.value) || tv.value.contains(value) {
                for label in &tv.labels {
                    labels.push(label);
                }
            }
        }
        labels
    }

    /// Propagate taint from one value to another.
    ///
    /// When data flows from one move's output to another move's input, this
    /// copies all taint labels from the source to the destination, extending
    /// each label's propagation chain. Like blood on the canvas spreading
    /// from one round to the next.
    pub fn propagate(&mut self, from: &str, to: &str) {
        // Collect labels from the source value.
        let mut propagated_labels: Vec<TaintLabel> = Vec::new();
        for tv in &self.values {
            if tv.value == from || from.contains(&tv.value) || tv.value.contains(from) {
                for label in &tv.labels {
                    let mut new_label = label.clone();
                    new_label
                        .propagation_chain
                        .push(format!("{} -> {}", from, to));
                    propagated_labels.push(new_label);
                }
            }
        }

        if propagated_labels.is_empty() {
            return;
        }

        // Apply labels to the destination value.
        for tv in &mut self.values {
            if tv.value == to {
                tv.labels.extend(propagated_labels);
                return;
            }
        }

        // Destination not yet tracked — create it.
        self.values.push(TaintedValue {
            value: to.to_string(),
            labels: propagated_labels,
        });
    }
}

// ---------------------------------------------------------------------------
// Shell bleed detection
// ---------------------------------------------------------------------------

/// A warning raised by the shell bleed detector — something leaked where it
/// shouldn't have, like blood between rounds.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShellBleedWarning {
    /// Name of the pattern that matched (e.g. "aws_access_key").
    pub pattern_name: String,
    /// The text that matched the pattern.
    pub matched_text: String,
    /// How severe the bleed is.
    pub severity: Sensitivity,
    /// Where in the command/environment the match was found.
    pub location: String,
}

/// Internal representation of a compiled secret pattern.
#[derive(Debug, Clone)]
struct SecretPattern {
    /// Human-readable name for this pattern.
    name: String,
    /// Compiled regex.
    regex: Regex,
    /// Severity if this pattern matches.
    severity: Sensitivity,
}

/// Scans shell commands and environment variables for leaked secrets before
/// a shell move lands.
///
/// The bleed detector watches for sensitive data escaping into shell commands
/// — API keys, tokens, private keys, and other secrets that should never
/// leave the ring unprotected.
#[derive(Debug, Clone)]
pub struct ShellBleedDetector {
    /// Registered secret patterns.
    patterns: Vec<SecretPattern>,
}

impl Default for ShellBleedDetector {
    fn default() -> Self {
        Self::new()
    }
}

impl ShellBleedDetector {
    /// Create a new bleed detector with built-in patterns for common secret
    /// types. Ready to scan from the first bell.
    pub fn new() -> Self {
        let mut detector = Self {
            patterns: Vec::new(),
        };
        detector.register_builtin_patterns();
        detector
    }

    /// Register built-in patterns for common secret formats.
    fn register_builtin_patterns(&mut self) {
        let builtins: &[(&str, &str, Sensitivity)] = &[
            // AWS access key IDs start with AKIA followed by 16 alphanumeric chars.
            ("aws_access_key", r"AKIA[0-9A-Z]{16}", Sensitivity::Secret),
            // Bearer tokens in command arguments.
            (
                "bearer_token",
                r"[Bb]earer\s+[A-Za-z0-9\-._~+/]+=*",
                Sensitivity::Secret,
            ),
            // Passwords embedded in URLs (https://user:pass@host).
            (
                "password_in_url",
                r"://[^/\s]+:[^@/\s]+@",
                Sensitivity::Confidential,
            ),
            // PEM private key markers.
            (
                "private_key",
                r"-----BEGIN\s+(RSA\s+)?PRIVATE\s+KEY-----",
                Sensitivity::Secret,
            ),
            // Generic API key patterns (key=..., api_key=..., apikey=...).
            (
                "generic_api_key",
                r#"(?i)(api[_\-]?key|api[_\-]?secret|access[_\-]?token|auth[_\-]?token)\s*[=:]\s*['\"]?[A-Za-z0-9\-._~+/]{16,}"#,
                Sensitivity::Confidential,
            ),
            // Long base64-encoded strings (40+ chars) that look like encoded secrets.
            (
                "base64_secret",
                r"[A-Za-z0-9+/]{40,}={0,2}",
                Sensitivity::Internal,
            ),
        ];

        for (name, pattern, severity) in builtins {
            if let Ok(regex) = Regex::new(pattern) {
                self.patterns.push(SecretPattern {
                    name: name.to_string(),
                    regex,
                    severity: *severity,
                });
            }
        }
    }

    /// Add a custom secret pattern to scan for.
    ///
    /// The pattern is a regex string. If the regex fails to compile, the
    /// pattern is silently ignored (no unwrap — we don't panic in the ring).
    pub fn add_secret_pattern(&mut self, name: &str, pattern: &str) {
        if let Ok(regex) = Regex::new(pattern) {
            self.patterns.push(SecretPattern {
                name: name.to_string(),
                regex,
                severity: Sensitivity::Confidential,
            });
        }
    }

    /// Scan a shell command string for leaked secrets.
    ///
    /// Returns a list of bleed warnings for every pattern that matched.
    /// If the command is clean, the vec is empty — no blood on the canvas.
    pub fn scan_command(&self, command: &str) -> Vec<ShellBleedWarning> {
        let mut warnings = Vec::new();

        for pattern in &self.patterns {
            for m in pattern.regex.find_iter(command) {
                warnings.push(ShellBleedWarning {
                    pattern_name: pattern.name.clone(),
                    matched_text: m.as_str().to_string(),
                    severity: pattern.severity,
                    location: "command".to_string(),
                });
            }
        }

        warnings
    }

    /// Scan environment variables that would be passed to a subprocess.
    ///
    /// Checks both variable names and values for secret patterns —
    /// bleed detection covers every angle of the ring.
    pub fn scan_environment(&self, env_vars: &[(String, String)]) -> Vec<ShellBleedWarning> {
        let mut warnings = Vec::new();

        for (key, value) in env_vars {
            let combined = format!("{}={}", key, value);
            for pattern in &self.patterns {
                for m in pattern.regex.find_iter(&combined) {
                    warnings.push(ShellBleedWarning {
                        pattern_name: pattern.name.clone(),
                        matched_text: m.as_str().to_string(),
                        severity: pattern.severity,
                        location: format!("env:{}", key),
                    });
                }
            }
        }

        warnings
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Build a fake AWS access key dynamically so static scanners don't flag
    /// the test source itself.
    fn fake_aws_key() -> String {
        // Prefix + 16 uppercase alphanumeric chars = valid AKIA pattern.
        let prefix = "AKIA";
        let suffix = "IOSFODNN7EXAMPLE";
        format!("{}{}", prefix, suffix)
    }

    /// Build a PEM private key header dynamically to avoid static scanner
    /// flags on the source code.
    fn pem_begin_header() -> String {
        let marker = "BEGIN RSA PRIVATE KEY";
        format!("-----{}-----", marker)
    }

    fn pem_end_header() -> String {
        let marker = "END RSA PRIVATE KEY";
        format!("-----{}-----", marker)
    }

    // -- TaintTracker tests -------------------------------------------------

    #[test]
    fn test_taint_basic_taint_and_check() {
        let mut tracker = TaintTracker::new();
        tracker.taint(
            "my-secret-value",
            TaintSource::UserInput,
            Sensitivity::Secret,
        );

        let labels = tracker.check_taint("my-secret-value");
        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].source, TaintSource::UserInput);
        assert_eq!(labels[0].sensitivity, Sensitivity::Secret);
    }

    #[test]
    fn test_taint_substring_detection() {
        let mut tracker = TaintTracker::new();
        tracker.taint(
            "API_KEY=abc123secret",
            TaintSource::Environment,
            Sensitivity::Secret,
        );

        // Checking a substring of the tainted value should match.
        let labels = tracker.check_taint("abc123secret");
        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].source, TaintSource::Environment);
    }

    #[test]
    fn test_taint_propagation_chain() {
        let mut tracker = TaintTracker::new();
        tracker.taint(
            "original-secret",
            TaintSource::UserInput,
            Sensitivity::Secret,
        );

        tracker.propagate("original-secret", "derived-value");

        let labels = tracker.check_taint("derived-value");
        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].propagation_chain.len(), 1);
        assert!(labels[0].propagation_chain[0].contains("original-secret"));
        assert!(labels[0].propagation_chain[0].contains("derived-value"));
    }

    #[test]
    fn test_taint_no_taint_on_clean_value() {
        let mut tracker = TaintTracker::new();
        tracker.taint("secret-data", TaintSource::UserInput, Sensitivity::Secret);

        let labels = tracker.check_taint("completely-unrelated");
        assert!(labels.is_empty());
    }

    #[test]
    fn test_taint_multiple_labels_on_same_value() {
        let mut tracker = TaintTracker::new();
        tracker.taint(
            "shared-value",
            TaintSource::UserInput,
            Sensitivity::Internal,
        );
        tracker.taint(
            "shared-value",
            TaintSource::ExternalApi("stripe".to_string()),
            Sensitivity::Secret,
        );

        let labels = tracker.check_taint("shared-value");
        assert_eq!(labels.len(), 2);

        let sources: Vec<&TaintSource> = labels.iter().map(|l| &l.source).collect();
        assert!(sources.contains(&&TaintSource::UserInput));
        assert!(sources.contains(&&TaintSource::ExternalApi("stripe".to_string())));
    }

    // -- ShellBleedDetector tests -------------------------------------------

    #[test]
    fn test_bleed_detect_aws_access_key() {
        let detector = ShellBleedDetector::new();
        let key = fake_aws_key();
        let cmd = format!("aws s3 cp --access-key {} s3://bucket", key);
        let warnings = detector.scan_command(&cmd);

        assert!(!warnings.is_empty());
        let aws_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "aws_access_key")
            .expect("should detect AWS access key");
        assert_eq!(aws_warning.severity, Sensitivity::Secret);
        assert!(aws_warning.matched_text.starts_with("AKIA"));
    }

    #[test]
    fn test_bleed_detect_bearer_token() {
        let detector = ShellBleedDetector::new();
        let warnings =
            detector.scan_command("curl -H 'Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.test'");

        assert!(!warnings.is_empty());
        let token_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "bearer_token")
            .expect("should detect bearer token");
        assert_eq!(token_warning.severity, Sensitivity::Secret);
    }

    #[test]
    fn test_bleed_detect_password_in_url() {
        let detector = ShellBleedDetector::new();
        let warnings = detector.scan_command("curl https://admin:supersecret@db.example.com/data");

        assert!(!warnings.is_empty());
        let pw_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "password_in_url")
            .expect("should detect password in URL");
        assert_eq!(pw_warning.severity, Sensitivity::Confidential);
    }

    #[test]
    fn test_bleed_detect_private_key() {
        let detector = ShellBleedDetector::new();
        let begin = pem_begin_header();
        let end = pem_end_header();
        let cmd = format!("echo '{}\nMIIEow...\n{}' > /tmp/key", begin, end);
        let warnings = detector.scan_command(&cmd);

        assert!(!warnings.is_empty());
        let key_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "private_key")
            .expect("should detect private key marker");
        assert_eq!(key_warning.severity, Sensitivity::Secret);
    }

    #[test]
    fn test_bleed_clean_command_passes() {
        let detector = ShellBleedDetector::new();
        let warnings = detector.scan_command("ls -la /tmp");

        assert!(
            warnings.is_empty(),
            "clean command should produce no warnings"
        );
    }

    #[test]
    fn test_bleed_detect_base64_encoded_secret() {
        let detector = ShellBleedDetector::new();
        // A long base64 string that looks like an encoded secret.
        let long_b64 = "YWJjZGVmZ2hpamtsbW5vcHFyc3R1dnd4eXoxMjM0NTY3ODk=";
        let cmd = format!("export SECRET={}", long_b64);
        let warnings = detector.scan_command(&cmd);

        assert!(!warnings.is_empty());
        let b64_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "base64_secret")
            .expect("should detect base64-encoded secret");
        assert_eq!(b64_warning.severity, Sensitivity::Internal);
    }

    #[test]
    fn test_bleed_env_var_scanning() {
        let detector = ShellBleedDetector::new();
        let key = fake_aws_key();
        let env_vars = vec![
            ("PATH".to_string(), "/usr/bin:/usr/local/bin".to_string()),
            ("AWS_ACCESS_KEY_ID".to_string(), key),
        ];

        let warnings = detector.scan_environment(&env_vars);
        assert!(!warnings.is_empty());

        let aws_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "aws_access_key")
            .expect("should detect AWS key in env vars");
        assert!(aws_warning.location.starts_with("env:"));
    }

    #[test]
    fn test_bleed_custom_pattern_registration() {
        let mut detector = ShellBleedDetector::new();
        detector.add_secret_pattern("github_token", r"ghp_[A-Za-z0-9]{36}");

        let token = format!("ghp_{}", "A".repeat(36));
        let cmd = format!("git clone https://{}@github.com/repo", token);
        let warnings = detector.scan_command(&cmd);

        assert!(!warnings.is_empty());
        let gh_warning = warnings
            .iter()
            .find(|w| w.pattern_name == "github_token")
            .expect("should detect custom GitHub token pattern");
        assert_eq!(gh_warning.severity, Sensitivity::Confidential);
    }

    #[test]
    fn test_sensitivity_ordering() {
        // Sensitivity should be ordered: Public < Internal < Confidential < Secret.
        assert!(Sensitivity::Public < Sensitivity::Internal);
        assert!(Sensitivity::Internal < Sensitivity::Confidential);
        assert!(Sensitivity::Confidential < Sensitivity::Secret);
        assert!(Sensitivity::Public < Sensitivity::Secret);
    }

    #[test]
    fn test_integration_shell_command_with_leaked_secret_blocked() {
        // Simulate the integration flow: taint a value, then scan a command
        // that contains it and verify the bleed detector catches it.
        let mut tracker = TaintTracker::new();
        let secret_key = fake_aws_key();
        tracker.taint(&secret_key, TaintSource::Environment, Sensitivity::Secret);

        let command = format!("aws s3 ls --access-key {}", secret_key);

        // Check taint on the command.
        let taint_labels = tracker.check_taint(&command);
        assert!(
            !taint_labels.is_empty(),
            "command containing tainted value should be detected"
        );

        // Also run the bleed detector.
        let detector = ShellBleedDetector::new();
        let warnings = detector.scan_command(&command);
        assert!(
            !warnings.is_empty(),
            "bleed detector should catch the AWS key"
        );

        // Verify that at least one warning is Secret severity.
        let has_secret = warnings
            .iter()
            .any(|w| w.severity >= Sensitivity::Confidential);
        assert!(
            has_secret,
            "leaked secret should produce at least Confidential severity"
        );
    }
}