parsentry_parser/

patterns.rs

//! Security pattern matching for vulnerability detection.

use parsentry_core::Language;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::Path;
use streaming_iterator::StreamingIterator;
use tree_sitter::{Language as TreeSitterLanguage, Parser, Query, QueryCursor};

/// Configuration for a security pattern.
#[derive(Debug, Clone, Deserialize)]
pub struct PatternConfig {
    #[serde(flatten)]
    pub pattern_type: PatternQuery,
    pub description: String,
    pub attack_vector: Vec<String>,
}

/// Query type for pattern matching.
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum PatternQuery {
    Definition { definition: String },
    Reference { reference: String },
}

/// Language-specific patterns configuration.
#[derive(Debug, Clone, Deserialize)]
pub struct LanguagePatterns {
    pub principals: Option<Vec<PatternConfig>>,
    pub actions: Option<Vec<PatternConfig>>,
    pub resources: Option<Vec<PatternConfig>>,
}

/// Security risk pattern matcher.
pub struct SecurityRiskPatterns {
    definition_queries: Vec<Query>,
    reference_queries: Vec<Query>,
    language: TreeSitterLanguage,
    pattern_configs: Vec<PatternConfig>,
}

/// A matched security pattern.
#[derive(Debug, Clone)]
pub struct PatternMatch {
    pub pattern_config: PatternConfig,
    pub start_byte: usize,
    pub end_byte: usize,
    pub matched_text: String,
}

impl SecurityRiskPatterns {
    /// Create a new pattern matcher for the given language.
    #[must_use]
    pub fn new(language: Language) -> Self {
        Self::new_with_root(language, None)
    }

    /// Create a new pattern matcher with a custom root directory for patterns.
    #[must_use]
    pub fn new_with_root(language: Language, root_dir: Option<&Path>) -> Self {
        let pattern_map = Self::load_patterns(root_dir);
        let lang_patterns = pattern_map
            .get(&language)
            .or_else(|| pattern_map.get(&Language::Other))
            .unwrap_or(&LanguagePatterns {
                principals: None,
                actions: None,
                resources: None,
            });

        let ts_language = Self::get_tree_sitter_language(language);

        let mut definition_queries = Vec::new();
        let mut reference_queries = Vec::new();
        let mut pattern_configs = Vec::new();

        // Collect all patterns from principals, actions, and resources into a flat list
        let all_configs: Vec<&PatternConfig> = lang_patterns
            .principals
            .iter()
            .chain(lang_patterns.actions.iter())
            .chain(lang_patterns.resources.iter())
            .flat_map(|v| v.iter())
            .collect();

        for config in all_configs {
            pattern_configs.push(config.clone());
            match &config.pattern_type {
                PatternQuery::Definition { definition } => {
                    if let Ok(query) = Query::new(&ts_language, definition) {
                        definition_queries.push(query);
                    }
                }
                PatternQuery::Reference { reference } => {
                    if let Ok(query) = Query::new(&ts_language, reference) {
                        reference_queries.push(query);
                    }
                }
            }
        }

        Self {
            definition_queries,
            reference_queries,
            language: ts_language,
            pattern_configs,
        }
    }

    fn get_tree_sitter_language(language: Language) -> TreeSitterLanguage {
        match language {
            Language::Python => tree_sitter_python::LANGUAGE.into(),
            Language::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
            Language::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
            Language::Rust => tree_sitter_rust::LANGUAGE.into(),
            Language::Java => tree_sitter_java::LANGUAGE.into(),
            Language::Go => tree_sitter_go::LANGUAGE.into(),
            Language::Ruby => tree_sitter_ruby::LANGUAGE.into(),
            Language::C => tree_sitter_c::LANGUAGE.into(),
            Language::Cpp => tree_sitter_cpp::LANGUAGE.into(),
            Language::Terraform => tree_sitter_hcl::LANGUAGE.into(),
            Language::Php => tree_sitter_php::LANGUAGE_PHP.into(),
            Language::Yaml => tree_sitter_yaml::LANGUAGE.into(),
            _ => tree_sitter_javascript::LANGUAGE.into(),
        }
    }

    /// Check if content matches any security pattern.
    #[must_use]
    pub fn matches(&self, content: &str) -> bool {
        let mut parser = Parser::new();
        if parser.set_language(&self.language).is_err() {
            return false;
        }

        let tree = match parser.parse(content, None) {
            Some(tree) => tree,
            None => return false,
        };

        let root_node = tree.root_node();

        let all_queries = [&self.definition_queries, &self.reference_queries];

        for query_set in all_queries {
            for query in query_set {
                let mut cursor = QueryCursor::new();
                let mut matches = cursor.matches(query, root_node, content.as_bytes());
                while let Some(match_) = matches.next() {
                    let mut has_valid_capture = false;

                    for capture in match_.captures {
                        let capture_name = &query.capture_names()[capture.index as usize];
                        let node = capture.node;
                        let start_byte = node.start_byte();
                        let end_byte = node.end_byte();
                        let matched_text = content[start_byte..end_byte].to_string();

                        if matched_text.trim().len() <= 2 {
                            continue;
                        }

                        match *capture_name {
                            "function" | "definition" | "class" | "method_def" | "call"
                            | "expression" | "attribute" => {
                                has_valid_capture = true;
                            }
                            "name" | "func" | "attr" | "obj" | "method" => {
                                has_valid_capture = true;
                            }
                            _ => {
                                has_valid_capture = true;
                            }
                        }
                    }

                    if has_valid_capture {
                        return true;
                    }
                }
            }
        }

        false
    }

    /// Get attack vectors for content.
    #[must_use]
    pub fn get_attack_vectors(&self, _content: &str) -> Vec<String> {
        Vec::new()
    }

    /// Get all pattern matches in content.
    #[must_use]
    pub fn get_pattern_matches(&self, content: &str) -> Vec<PatternMatch> {
        let mut parser = Parser::new();
        if parser.set_language(&self.language).is_err() {
            return Vec::new();
        }

        let tree = match parser.parse(content, None) {
            Some(tree) => tree,
            None => return Vec::new(),
        };

        let root_node = tree.root_node();
        let mut pattern_matches = Vec::new();
        let content_bytes = content.as_bytes();

        let mut process_queries = |queries: &[Query], is_definition: bool| {
            for (query_idx, query) in queries.iter().enumerate() {
                let mut cursor = QueryCursor::new();
                let mut matches = cursor.matches(query, root_node, content_bytes);

                while let Some(match_) = matches.next() {
                    let mut best_node = None;
                    let mut best_text = String::new();
                    let mut best_priority = 0;

                    for capture in match_.captures {
                        let capture_name = &query.capture_names()[capture.index as usize];
                        let node = capture.node;
                        let start_byte = node.start_byte();
                        let end_byte = node.end_byte();
                        let matched_text = content[start_byte..end_byte].to_string();

                        if matched_text.trim().len() <= 2 {
                            continue;
                        }

                        let (priority, candidate_node, candidate_text) = match *capture_name {
                            "function" | "definition" | "class" | "method_def" => {
                                (100, Some(node), matched_text.clone())
                            }
                            "call" | "expression" | "attribute" => {
                                (90, Some(node), matched_text.clone())
                            }
                            "name" | "func" | "attr" | "obj" | "method" => {
                                let mut found_parent = None;
                                let mut parent = node.parent();
                                while let Some(p) = parent {
                                    if (is_definition
                                        && (p.kind().contains("definition")
                                            || p.kind().contains("declaration")))
                                        || (!is_definition
                                            && (p.kind().contains("call")
                                                || p.kind().contains("attribute")
                                                || p.kind().contains("expression")))
                                    {
                                        found_parent = Some(p);
                                        break;
                                    }
                                    parent = p.parent();
                                }
                                if let Some(p) = found_parent {
                                    (
                                        80,
                                        Some(p),
                                        content[p.start_byte()..p.end_byte()].to_string(),
                                    )
                                } else {
                                    (70, Some(node), matched_text.clone())
                                }
                            }
                            "param" | "func_name" => {
                                let mut found_func = None;
                                let mut parent = node.parent();
                                while let Some(p) = parent {
                                    if p.kind() == "function_definition" {
                                        found_func = Some(p);
                                        break;
                                    }
                                    parent = p.parent();
                                }
                                if let Some(p) = found_func {
                                    (
                                        85,
                                        Some(p),
                                        content[p.start_byte()..p.end_byte()].to_string(),
                                    )
                                } else {
                                    (60, Some(node), matched_text.clone())
                                }
                            }
                            _ => (50, Some(node), matched_text.clone()),
                        };

                        if priority > best_priority {
                            best_priority = priority;
                            best_node = candidate_node;
                            best_text = candidate_text;
                        }
                    }

                    if let Some(node) = best_node {
                        let start_byte = node.start_byte();
                        let end_byte = node.end_byte();

                        // Find the matching config by counting definition/reference queries
                        let mut config_idx = 0;
                        for config in &self.pattern_configs {
                            let matches_type = matches!(
                                (&config.pattern_type, is_definition),
                                (PatternQuery::Definition { .. }, true)
                                    | (PatternQuery::Reference { .. }, false)
                            );

                            if matches_type {
                                if config_idx == query_idx {
                                    pattern_matches.push(PatternMatch {
                                        pattern_config: config.clone(),
                                        start_byte,
                                        end_byte,
                                        matched_text: best_text.clone(),
                                    });
                                    break;
                                }
                                config_idx += 1;
                            }
                        }
                    }
                }
            }
        };

        process_queries(&self.definition_queries, true);
        process_queries(&self.reference_queries, false);

        pattern_matches
    }

    fn load_patterns(root_dir: Option<&Path>) -> HashMap<Language, LanguagePatterns> {
        use Language::*;

        let mut map = HashMap::new();

        let languages = [
            (Python, include_str!("patterns/python.yml")),
            (JavaScript, include_str!("patterns/javascript.yml")),
            (Rust, include_str!("patterns/rust.yml")),
            (TypeScript, include_str!("patterns/typescript.yml")),
            (Java, include_str!("patterns/java.yml")),
            (Go, include_str!("patterns/go.yml")),
            (Ruby, include_str!("patterns/ruby.yml")),
            (C, include_str!("patterns/c.yml")),
            (Cpp, include_str!("patterns/cpp.yml")),
            (Php, include_str!("patterns/php.yml")),
            (Terraform, include_str!("patterns/terraform.yml")),
        ];

        for (lang, content) in languages {
            match serde_yaml::from_str::<LanguagePatterns>(content) {
                Ok(patterns) => {
                    map.insert(lang, patterns);
                }
                Err(e) => {
                    eprintln!("Failed to parse patterns for {:?}: {}", lang, e);
                }
            }
        }

        // Load CI/CD platform patterns and merge into Yaml language
        let cicd_patterns = [
            include_str!("patterns/github-actions.yml"), // GitHub Actions
            include_str!("patterns/gitlab-ci.yml"),
            include_str!("patterns/circleci.yml"),
            include_str!("patterns/travis.yml"),
            include_str!("patterns/jenkins.yml"),
        ];

        let mut merged_yaml_patterns = LanguagePatterns {
            principals: Some(Vec::new()),
            actions: Some(Vec::new()),
            resources: Some(Vec::new()),
        };

        for content in cicd_patterns {
            if let Ok(patterns) = serde_yaml::from_str::<LanguagePatterns>(content) {
                if let Some(principals) = patterns.principals {
                    merged_yaml_patterns
                        .principals
                        .as_mut()
                        .unwrap()
                        .extend(principals);
                }
                if let Some(actions) = patterns.actions {
                    merged_yaml_patterns
                        .actions
                        .as_mut()
                        .unwrap()
                        .extend(actions);
                }
                if let Some(resources) = patterns.resources {
                    merged_yaml_patterns
                        .resources
                        .as_mut()
                        .unwrap()
                        .extend(resources);
                }
            }
        }

        map.insert(Yaml, merged_yaml_patterns);

        Self::load_custom_patterns(&mut map, root_dir);

        map
    }

    /// Add dynamic queries (e.g. from threat model) at runtime.
    /// `query_type` is "definition" or "reference".
    pub fn add_query(
        &mut self,
        query_type: &str,
        query_str: &str,
        description: &str,
        attack_vector: Vec<String>,
    ) -> bool {
        let query = match Query::new(&self.language, query_str) {
            Ok(q) => q,
            Err(_) => return false,
        };

        let is_definition = query_type == "definition";
        let pattern_query = if is_definition {
            PatternQuery::Definition {
                definition: query_str.to_string(),
            }
        } else {
            PatternQuery::Reference {
                reference: query_str.to_string(),
            }
        };

        let config = PatternConfig {
            pattern_type: pattern_query,
            description: description.to_string(),
            attack_vector,
        };

        self.pattern_configs.push(config);

        if is_definition {
            self.definition_queries.push(query);
        } else {
            self.reference_queries.push(query);
        }

        true
    }

    fn load_custom_patterns(
        map: &mut HashMap<Language, LanguagePatterns>,
        root_dir: Option<&Path>,
    ) {
        let vuln_patterns_path = if let Some(root) = root_dir {
            root.join("vuln-patterns.yml")
        } else {
            Path::new("vuln-patterns.yml").to_path_buf()
        };

        if vuln_patterns_path.exists() {
            match std::fs::read_to_string(&vuln_patterns_path) {
                Ok(content) => {
                    match serde_yaml::from_str::<HashMap<String, LanguagePatterns>>(&content) {
                        Ok(custom_patterns) => {
                            for (lang_name, patterns) in custom_patterns {
                                let language = match lang_name.as_str() {
                                    "Python" => Language::Python,
                                    "JavaScript" => Language::JavaScript,
                                    "TypeScript" => Language::TypeScript,
                                    "Rust" => Language::Rust,
                                    "Java" => Language::Java,
                                    "Go" => Language::Go,
                                    "Ruby" => Language::Ruby,
                                    "C" => Language::C,
                                    "Cpp" => Language::Cpp,
                                    "Terraform" => Language::Terraform,
                                    "CloudFormation" => Language::CloudFormation,
                                    "Kubernetes" => Language::Kubernetes,
                                    "YAML" => Language::Yaml,
                                    "GitLabCI" => Language::Yaml,
                                    "CircleCI" => Language::Yaml,
                                    "TravisCI" => Language::Yaml,
                                    "Jenkins" => Language::Yaml,
                                    "Bash" => Language::Bash,
                                    "Shell" => Language::Shell,
                                    "Php" | "PHP" => Language::Php,
                                    _ => continue,
                                };

                                match map.get_mut(&language) {
                                    Some(existing) => {
                                        if let Some(custom_principals) = patterns.principals {
                                            match &mut existing.principals {
                                                Some(principals) => {
                                                    principals.extend(custom_principals)
                                                }
                                                None => {
                                                    existing.principals = Some(custom_principals)
                                                }
                                            }
                                        }
                                        if let Some(custom_actions) = patterns.actions {
                                            match &mut existing.actions {
                                                Some(actions) => actions.extend(custom_actions),
                                                None => existing.actions = Some(custom_actions),
                                            }
                                        }
                                        if let Some(custom_resources) = patterns.resources {
                                            match &mut existing.resources {
                                                Some(resources) => {
                                                    resources.extend(custom_resources)
                                                }
                                                None => existing.resources = Some(custom_resources),
                                            }
                                        }
                                    }
                                    None => {
                                        map.insert(language, patterns);
                                    }
                                }
                            }
                        }
                        Err(e) => {
                            eprintln!("Failed to parse vuln-patterns.yml: {}", e);
                        }
                    }
                }
                Err(e) => {
                    eprintln!("Failed to read vuln-patterns.yml: {}", e);
                }
            }
        }
    }
}