perimeterx-fastly-enforcer 2.2.2

use crate::px_debug;
use crate::pxconfig::PXConfig;
use apollo_parser::{
    Parser,
    cst::{self, CstNode},
};
use regex::Regex;
use serde_json;

const GRAPHQL_KEYWORD_MAX_LEN: usize = 100;
const GRAPHQL_KEYWORDS_MAX_COUNT: usize = 500;

#[derive(Default)]
pub struct PXGraphQLExtractedItem {
    pub op_type: String,
    pub name: String,
    pub sensitive: bool,
    pub variables: Vec<String>,
    pub keywords: Vec<String>,
}

// GraphQL operation names are strictly case-sensitive
fn is_sensitive_graphql_operation_name(names: &[String], name: &str) -> bool {
    let name = name.trim();
    if name.is_empty() {
        return false;
    }
    names.iter().any(|n| n.trim() == name)
}

fn clean_keyword(kw: &str) -> String {
    let trimmed = kw.trim();
    trimmed.chars().take(GRAPHQL_KEYWORD_MAX_LEN).collect()
}

fn extract_query_keywords(query: &str, patterns: &[Regex]) -> Vec<String> {
    if patterns.is_empty() {
        return Vec::new();
    }

    let mut keywords = Vec::new();
    for pattern in patterns {
        for m in pattern.find_iter(query) {
            let kw = clean_keyword(m.as_str());
            if kw.is_empty() {
                continue;
            }
            keywords.push(kw);
            if keywords.len() >= GRAPHQL_KEYWORDS_MAX_COUNT {
                return keywords;
            }
        }
    }
    keywords
}

// Extract variable key paths for leaf values only; object values are traversed, not reported.
fn traverse_json_for_keys(value: &serde_json::Value, keys: &mut Vec<String>, current_key: &str) {
    let Some(map) = value.as_object() else {
        return;
    };

    for (key, val) in map {
        let mut new_key = current_key.to_owned();
        if !new_key.is_empty() {
            new_key.push('.');
        }
        new_key.push_str(key);

        if val.is_object() {
            traverse_json_for_keys(val, keys, &new_key);
        } else {
            keys.push(new_key);
        }
    }
}

// Extract GraphQL query from JSON value if it exists, otherwise return the JSON value as a string (for Automated Persisted Query APQ type of queries)
fn graphql_query_from_value(v: &serde_json::Value) -> String {
    if let Some(obj) = v.as_object() {
        match obj.get("query") {
            Some(query_val) => query_val.as_str().unwrap_or_default().to_string(),
            None => v.to_string(),
        }
    } else {
        v.as_str().unwrap_or_default().to_string()
    }
}

pub(crate) fn graphql_parse_and_extract(
    v: &serde_json::Value,
    conf: &PXConfig,
) -> Option<PXGraphQLExtractedItem> {
    let query = graphql_query_from_value(v);

    let mut item: PXGraphQLExtractedItem = Default::default();

    if let Some(var_obj) = v.get("variables") {
        if !var_obj.is_null() {
            traverse_json_for_keys(var_obj, &mut item.variables, "");
        }
    }

    let parser = Parser::new(query.as_str());
    let cst = parser.parse();

    let doc = cst.document();

    // find Operation, extract name and type
    for def in doc.definitions() {
        if let cst::Definition::OperationDefinition(op) = def {
            if let Some(n) = op.name() {
                item.name = n.text().to_string();
                if is_sensitive_graphql_operation_name(
                    &conf.sensitive_graphql_operation_names,
                    &item.name,
                ) {
                    item.sensitive = true;
                }
            };

            if let Some(t) = op.operation_type() {
                item.op_type = t.source_string().to_string();
            } else {
                item.op_type = "query".to_string();
            }
            if conf
                .sensitive_graphql_operation_types
                .contains(&item.op_type)
            {
                item.sensitive = true;
            }
        }
    }

    // for Automated Persisted Query (APQ) type of queries
    if item.name.is_empty() {
        if let Some(op_name) = v.get("operationName").and_then(|n| n.as_str()) {
            item.name = op_name.to_string();
            if is_sensitive_graphql_operation_name(
                &conf.sensitive_graphql_operation_names,
                &item.name,
            ) {
                item.sensitive = true;
            }
        }
    }

    // extract keywords
    if !conf.graphql_keywords.is_empty() && !query.is_empty() {
        item.keywords = extract_query_keywords(&query, &conf.graphql_keywords);
    }

    // enable sensitivity based on keywords
    if !item.sensitive
        && !conf.sensitive_graphql_operation_names.is_empty()
        && !item.keywords.is_empty()
    {
        item.sensitive = true;
    }

    Some(item)
}

// Extract GraphQL items from JSON body,
// Return a list of PXGraphQLExtractedItem
pub fn px_graphql_extract(v: &serde_json::Value, conf: &PXConfig) -> Vec<PXGraphQLExtractedItem> {
    let mut extracted_items = Vec::new();

    px_debug!("Extracting GraphQL items from JSON body");

    // GraphQL JSON could contain multiple objects
    if let Some(array) = v.as_array() {
        for i in array {
            if let Some(item) = graphql_parse_and_extract(i, conf) {
                extracted_items.push(item);
            }
        }
    } else if let Some(item) = graphql_parse_and_extract(v, conf) {
        extracted_items.push(item);
    }

    extracted_items
}

// return true if GraphQL request is "sensitive"
pub fn px_graphql_is_sensitive(items: &Vec<PXGraphQLExtractedItem>) -> bool {
    for item in items {
        if item.sensitive {
            return true;
        }
    }
    false
}