thoughtjack 0.6.0

//! Built-in attack scenarios (TJ-SPEC-010)
//!
//! OATF scenarios auto-discovered at build time from `scenarios/library/`.
//! Only documents that pass `oatf::load()` validation are embedded.
//! Run `scenarios/fetch.sh` to download the OATF official scenario library.

use std::fmt;

// ============================================================================
// Types
// ============================================================================

/// A built-in scenario embedded in the binary.
///
/// Metadata is extracted from the OATF document at build time.
/// The YAML content is included verbatim via `include_str!`.
///
/// Implements: TJ-SPEC-010 F-001
pub struct BuiltinScenario {
    /// Unique identifier derived from `attack.id` (lowercase, e.g., "oatf-001").
    pub name: &'static str,

    /// Short human-readable description from `attack.name`.
    pub description: &'static str,

    /// Category for organization, mapped from `attack.classification.category`.
    pub category: ScenarioCategory,

    /// Tags from `attack.classification.tags`.
    pub tags: &'static [&'static str],

    /// Raw YAML content (embedded at compile time).
    pub yaml: &'static str,
}

/// Category for organizing built-in scenarios.
///
/// Mapped from OATF `classification.category` values at build time.
///
/// Implements: TJ-SPEC-010 F-001
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum ScenarioCategory {
    /// Prompt injection, capability poisoning, and content manipulation.
    Injection,
    /// Denial of service and availability disruption.
    #[value(name = "dos")]
    DoS,
    /// Temporal attacks (rug pulls, sleepers).
    Temporal,
    /// Resource and subscription attacks.
    Resource,
    /// Protocol-level attacks and oversight bypass.
    Protocol,
    /// Compound attacks combining multiple techniques.
    #[value(name = "multi_vector")]
    MultiVector,
}

impl ScenarioCategory {
    /// Returns the human-readable title-case label.
    ///
    /// Implements: TJ-SPEC-010 F-004
    #[must_use]
    pub const fn label(self) -> &'static str {
        match self {
            Self::Injection => "Injection",
            Self::DoS => "DoS",
            Self::Temporal => "Temporal",
            Self::Resource => "Resource",
            Self::Protocol => "Protocol",
            Self::MultiVector => "Multi-Vector",
        }
    }

    /// Returns all category variants in display order.
    #[must_use]
    pub const fn all() -> &'static [Self] {
        &[
            Self::Temporal,
            Self::Injection,
            Self::Resource,
            Self::DoS,
            Self::Protocol,
            Self::MultiVector,
        ]
    }
}

impl fmt::Display for ScenarioCategory {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Injection => write!(f, "injection"),
            Self::DoS => write!(f, "dos"),
            Self::Temporal => write!(f, "temporal"),
            Self::Resource => write!(f, "resource"),
            Self::Protocol => write!(f, "protocol"),
            Self::MultiVector => write!(f, "multi_vector"),
        }
    }
}

// ============================================================================
// Registry (auto-generated by build.rs)
// ============================================================================

/// Global registry of all built-in scenarios, discovered at build time.
///
/// Implements: TJ-SPEC-010 F-001, F-002
static BUILTIN_SCENARIOS: &[BuiltinScenario] =
    include!(concat!(env!("OUT_DIR"), "/builtin_scenarios.rs"));

// ============================================================================
// Public API
// ============================================================================

/// Look up a scenario by exact name.
///
/// Implements: TJ-SPEC-010 F-001
#[must_use]
pub fn find_scenario(name: &str) -> Option<&'static BuiltinScenario> {
    BUILTIN_SCENARIOS.iter().find(|s| s.name == name)
}

/// List all scenarios, optionally filtered by category and/or tag.
///
/// Implements: TJ-SPEC-010 F-004
#[must_use]
pub fn list_scenarios(
    category: Option<ScenarioCategory>,
    tag: Option<&str>,
) -> Vec<&'static BuiltinScenario> {
    BUILTIN_SCENARIOS
        .iter()
        .filter(|s| category.is_none_or(|c| s.category == c))
        .filter(|s| tag.is_none_or(|t| s.tags.contains(&t)))
        .collect()
}

/// Suggest a similar scenario name for typo correction.
///
/// Returns the closest match if its Damerau-Levenshtein distance is ≤ 3.
///
/// Implements: TJ-SPEC-010 F-006
#[must_use]
pub fn suggest_scenario(input: &str) -> Option<String> {
    BUILTIN_SCENARIOS
        .iter()
        .map(|s| (s.name, strsim::damerau_levenshtein(input, s.name)))
        .filter(|(_, dist)| *dist <= 3)
        .min_by_key(|(_, dist)| *dist)
        .map(|(name, _)| name.to_string())
}

/// Returns all scenario names in registry order.
///
/// Implements: TJ-SPEC-010 F-003
#[must_use]
pub fn list_scenario_names() -> Vec<&'static str> {
    BUILTIN_SCENARIOS.iter().map(|s| s.name).collect()
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn no_duplicate_scenario_names() {
        let names: Vec<&str> = list_scenarios(None, None).iter().map(|s| s.name).collect();
        let unique: HashSet<&str> = names.iter().copied().collect();
        assert_eq!(names.len(), unique.len(), "Duplicate scenario names found");
    }

    #[test]
    fn all_builtin_scenarios_are_self_contained() {
        for scenario in list_scenarios(None, None) {
            assert!(
                !scenario.yaml.contains("$include:"),
                "Built-in scenario '{}' contains $include directive",
                scenario.name,
            );
            assert!(
                !scenario.yaml.contains("$file:"),
                "Built-in scenario '{}' contains $file directive",
                scenario.name,
            );
            assert!(
                !scenario.yaml.contains("${env."),
                "Built-in scenario '{}' references environment variables",
                scenario.name,
            );
        }
    }

    #[test]
    fn builtin_scenarios_within_binary_size_budget() {
        let total_bytes: usize = list_scenarios(None, None)
            .iter()
            .map(|s| s.yaml.len())
            .sum();
        assert!(
            total_bytes < 1_000_000,
            "Total embedded YAML is {total_bytes} bytes, exceeds 1MB budget"
        );
    }

    #[test]
    fn find_scenario_missing() {
        assert!(find_scenario("nonexistent").is_none());
    }

    #[test]
    fn suggest_scenario_far() {
        // "xyzabc123" is too far from any scenario name
        let suggestion = suggest_scenario("xyzabc123");
        assert!(suggestion.is_none());
    }

    #[test]
    fn all_scenarios_pass_oatf_validation() {
        for scenario in list_scenarios(None, None) {
            let result = oatf::load(scenario.yaml);
            assert!(
                result.is_ok(),
                "Built-in scenario '{}' failed OATF validation: {:?}",
                scenario.name,
                result.err()
            );
        }
    }

    #[test]
    fn scenario_metadata_populated() {
        for scenario in list_scenarios(None, None) {
            assert!(!scenario.name.is_empty(), "Scenario name is empty");
            assert!(
                !scenario.description.is_empty(),
                "Scenario '{}' has empty description",
                scenario.name
            );
            assert!(
                !scenario.yaml.is_empty(),
                "Scenario '{}' has empty YAML",
                scenario.name
            );
        }
    }

    #[test]
    fn category_display_lowercase() {
        assert_eq!(ScenarioCategory::Injection.to_string(), "injection");
        assert_eq!(ScenarioCategory::DoS.to_string(), "dos");
        assert_eq!(ScenarioCategory::Temporal.to_string(), "temporal");
        assert_eq!(ScenarioCategory::Resource.to_string(), "resource");
        assert_eq!(ScenarioCategory::Protocol.to_string(), "protocol");
        assert_eq!(ScenarioCategory::MultiVector.to_string(), "multi_vector");
    }

    #[test]
    fn category_label_titlecase() {
        assert_eq!(ScenarioCategory::Injection.label(), "Injection");
        assert_eq!(ScenarioCategory::DoS.label(), "DoS");
        assert_eq!(ScenarioCategory::Temporal.label(), "Temporal");
        assert_eq!(ScenarioCategory::Resource.label(), "Resource");
        assert_eq!(ScenarioCategory::Protocol.label(), "Protocol");
        assert_eq!(ScenarioCategory::MultiVector.label(), "Multi-Vector");
    }
}