repotoire 0.8.0

//! AST-driven extraction of [`super::predict::Evidence`] for Python
//! pymongo / motor NoSQL-query call sites.
//!
//! # Why a separate module
//!
//! The scorer in [`super::predict`] takes plain data
//! ([`super::predict::Evidence`]) so it can be unit-tested without an
//! AST. This module's job is to populate that struct from a
//! `tree_sitter::Node` for a Python `call` expression whose callee
//! names a pymongo query API.
//!
//! Splitting the two halves matches Phase 2a/2b/2c/2d/2e/2f/2g/2h's
//! `evidence.rs` split.
//!
//! # What this module knows about
//!
//! - Walking the module AST to collect every call whose callee chain
//!   names a recognized pymongo / motor query method
//!   ([`collect_python_nosql_sites`]).
//! - Walking up from the call node to the enclosing
//!   `function_definition` (for name) and `class_definition`
//!   (informational).
//! - Detecting the function's decorator list (for the route-handler
//!   signal) and naming-convention match.
//! - **Structural classification of the query argument** — the heart
//!   of Phase 2i:
//!   * `**user_input` (dictionary_splat) in the query dict literal
//!     → `NosqlApi::DictExpansion` (D1.c collapse).
//!   * `$where` / `$function` / `$expr` / `$accumulator` key with a
//!     user-input expression value → `NosqlApi::OperatorInjection`
//!     (D1.b collapse).
//!   * All user-input values cast (`str`, `ObjectId`, `int`,
//!     pydantic), no dangerous operators, no expansion →
//!     `NosqlApi::TypedValueQuery` (D1.a collapse).
//!   * Otherwise → `NosqlApi::Ambiguous` (weighted scoring).
//! - Detecting `$ne` / `$gt` / `$lt` with a literal value (the
//!   developer-written operator FP-reduction signal).
//! - Detecting user-input source family (TypedString vs
//!   UnstructuredJson) within ±10 lines of the call.
//! - Detecting type-narrowing casts (`ObjectId`, `int`, pydantic)
//!   within ±5 lines of the call.
//! - Reading the source line for `# repotoire: nosql-safe[...]` /
//!   `nosql-vulnerable[...]` annotations.
//!
//! # What this module deliberately does NOT do
//!
//! - Does not look for evidence in non-Python languages (D5.1 scope —
//!   JS Mongoose deferred to a follow-on phase).
//! - Does not trace cross-statement query assembly
//!   (`q = build_query(req.json); users.find(q)`) — D5.3 v0 limitation;
//!   user annotates with `# repotoire: nosql-vulnerable[...]`.
//! - Does not trace pydantic validation across function boundaries —
//!   user annotates with `# repotoire: nosql-safe[pydantic-validated]`.
//! - Does not detect Redis `.eval`, Postgres `pg.query`, or other
//!   non-MongoDB query APIs (D5.5 scope).
//!
//! # Status
//!
//! Wired in via `mod.rs::scan_python_file_dual_branch` in Phase 2i's
//! commit 5 (integration). Every public symbol below is reachable
//! from the integration path.

use super::predict::{
    classify_user_input_source, extract_nosql_safe_reason, extract_nosql_vulnerable_source,
    line_contains_dangerous_operator, line_contains_developer_operator, line_contains_type_cast,
    matches_route_handler_decorator, matches_route_handler_name, matches_trust_boundary_name,
    Evidence, NosqlApi, UserInputSource, UNSTRUCTURED_JSON_USER_INPUT_SUBSTRINGS,
};
use crate::detectors::security::ast_helpers::{enclosing_python_function, node_text};
use tree_sitter::Node;

/// pymongo / motor query method names. The detection trigger is "call
/// expression whose callee chain ends with one of these names" — a
/// receiver-method-call shape (`coll.find(...)`, `db.users.find_one(...)`).
///
/// Mirrors the legacy `NOSQL_PATTERN` regex but encoded as a tree-
/// sitter-aware list so we can inspect the call's AST argument
/// structure rather than the line text alone.
pub(super) const PYMONGO_QUERY_METHODS: &[&str] = &[
    "find",
    "find_one",
    "find_by_id",
    "find_one_and_update",
    "find_one_and_delete",
    "find_one_and_replace",
    "update",
    "update_one",
    "update_many",
    "replace_one",
    "delete",
    "delete_one",
    "delete_many",
    "aggregate",
    "count",
    "count_documents",
    "estimated_document_count",
    "distinct",
];

/// A Python pymongo query call site discovered by walking the module
/// AST. Returned by [`collect_python_nosql_sites`] so the integration
/// in `mod.rs` can iterate without re-walking.
pub(super) struct PythonNosqlSite<'a> {
    pub call_node: Node<'a>,
    pub api: NosqlApi,
    /// The raw callee text used for the title/description (e.g.
    /// `users.find_one`, `db.users.aggregate`). Resolved from the call
    /// expression's `function` child.
    pub callee_label: String,
}

/// Walk a Python module AST and collect every call whose callee
/// receiver-chain ends with a pymongo query method name.
///
/// Performs structural API classification on the first positional
/// argument:
///
/// - `**user_input` in the query dict literal → `DictExpansion` (D1.c).
/// - Dangerous operator (`$where`/`$function`/`$expr`/`$accumulator`)
///   with user-input value → `OperatorInjection` (D1.b).
/// - All values cast or literal, no dangerous operators → `TypedValueQuery`
///   (D1.a).
/// - Otherwise → `Ambiguous` (weighted scoring).
///
/// Skips:
/// - Calls whose receiver chain doesn't end with a recognized pymongo
///   method name.
/// - Array-method receivers (`items.find(...)` where the receiver name
///   suggests a list, not a MongoDB collection) — preserves the
///   legacy `is_array_method` FP filter.
pub(super) fn collect_python_nosql_sites<'a>(
    module_root: Node<'a>,
    source: &'a [u8],
) -> Vec<PythonNosqlSite<'a>> {
    let mut sites = Vec::new();
    let mut stack: Vec<Node<'_>> = vec![module_root];
    while let Some(node) = stack.pop() {
        let mut cursor = node.walk();
        for child in node.children(&mut cursor) {
            stack.push(child);
        }
        if node.kind() != "call" {
            continue;
        }
        let Some(func) = node.child_by_field_name("function") else {
            continue;
        };
        let func_text = node_text(func, source).unwrap_or("");
        if func_text.is_empty() {
            continue;
        }

        // Recognized pymongo method-chain shape?
        if !looks_like_pymongo_call(func_text) {
            continue;
        }

        // Preserve the legacy `is_array_method` FP filter: skip calls
        // where the receiver name suggests a Python list, not a MongoDB
        // collection.
        if receiver_looks_like_array(func_text) {
            continue;
        }

        let api = classify_query_shape(node, source);

        sites.push(PythonNosqlSite {
            call_node: node,
            api,
            callee_label: func_text.to_string(),
        });
    }
    sites
}

/// True iff the callee text matches a `<receiver>.<pymongo_method>(...)`
/// shape. The receiver is anything (we trust the method-name lexicon).
///
/// Examples that match: `users.find`, `db.users.find_one`,
/// `mongo.users.aggregate`. Examples that don't: `request.find` (no —
/// but the array-method filter catches further FPs), bare `find(...)`.
fn looks_like_pymongo_call(func_text: &str) -> bool {
    // Receiver-method shape required (contains a dot).
    let Some(method_name) = func_text.rsplit('.').next() else {
        return false;
    };
    if method_name == func_text {
        // No dot — bare call. Not a method call.
        return false;
    }
    PYMONGO_QUERY_METHODS.contains(&method_name)
}

/// True iff the receiver name suggests this is a Python list/tuple's
/// `.find()` method (which doesn't exist for lists, but the user might
/// be calling `.index()` or `Array.find()` translated from JS). Mirrors
/// the legacy `is_array_method` filter.
///
/// Note: Python lists don't have `.find()` — they have `.index()`.
/// This filter is mostly defensive against:
/// - JS-style `items.find(predicate)` patterns inadvertently parsed as
///   Python (rare; tree-sitter-python would error).
/// - String `.find()` calls on identifiers named `items`/`list`/etc.
/// - Helper-function names that happen to end in `.find`.
fn receiver_looks_like_array(func_text: &str) -> bool {
    // Split into receiver and method.
    let parts: Vec<&str> = func_text.rsplitn(2, '.').collect();
    if parts.len() != 2 {
        return false;
    }
    let receiver = parts[1].to_lowercase();
    const ARRAY_RECEIVER_SUBSTRINGS: &[&str] = &[
        "items.", "items", "list", "array", "results", "options", "elements", "entries", "records",
        "rows", "values", "keys",
    ];
    // Match against the bare last-receiver segment (e.g. for `obj.items.find`,
    // the last receiver segment is `items`).
    let last_receiver = receiver.rsplit('.').next().unwrap_or(&receiver);
    ARRAY_RECEIVER_SUBSTRINGS
        .iter()
        .any(|s| last_receiver == *s || last_receiver == s.trim_end_matches('.'))
}

/// Structural classification of the query argument. The heart of
/// Phase 2i's evidence extraction.
///
/// Inspection priority:
/// 1. **`**dictionary_splat` of recognized user-input identifier** in
///    the first positional argument → `DictExpansion`.
/// 2. **Dangerous operator key + user-input value** → `OperatorInjection`.
/// 3. **All values cast or literal, no dangerous operators, no `**`** →
///    `TypedValueQuery`.
/// 4. Otherwise → `Ambiguous`.
fn classify_query_shape(call_node: Node<'_>, source: &[u8]) -> NosqlApi {
    let Some(args) = call_node.child_by_field_name("arguments") else {
        return NosqlApi::Ambiguous;
    };
    let Some(first_arg) = first_positional_arg(args) else {
        return NosqlApi::Ambiguous;
    };

    // ── Priority 1: Dict expansion of raw user input. ──
    if dict_has_user_input_splat(first_arg, source) {
        return NosqlApi::DictExpansion;
    }

    // For the remaining checks the arg must be a dict literal (or a
    // list of dicts for aggregate pipelines). Other shapes (identifier,
    // function call result) are Ambiguous — the predictor falls back
    // to weighted scoring.
    let dict_text = node_text(first_arg, source).unwrap_or("");

    // ── Priority 2: Dangerous operator with user-input value. ──
    if has_dangerous_operator_with_user_input(first_arg, source, dict_text) {
        return NosqlApi::OperatorInjection;
    }

    // ── Priority 3: Typed-value query (D1.a). ──
    // Conditions:
    //  - First arg is a dictionary literal.
    //  - No dangerous operators anywhere in the dict text.
    //  - No `**` expansion (already filtered above).
    //  - Every user-input value position is wrapped in a type cast
    //    OR there is no user input in the dict at all.
    if first_arg.kind() == "dictionary"
        && !line_contains_dangerous_operator(dict_text)
        && is_typed_value_query(first_arg, source)
    {
        return NosqlApi::TypedValueQuery;
    }

    NosqlApi::Ambiguous
}

/// Return the first positional argument node from an `argument_list`,
/// skipping `keyword_argument` children.
fn first_positional_arg<'a>(args: Node<'a>) -> Option<Node<'a>> {
    let mut cursor = args.walk();
    for child in args.children(&mut cursor) {
        if !child.is_named() {
            continue;
        }
        if child.kind() == "keyword_argument" {
            continue;
        }
        return Some(child);
    }
    None
}

/// True iff the argument is a dict literal containing
/// `**<user_input_identifier>` (e.g. `{**request.json}`,
/// `{**req.body, "other": 1}`). The presence of `**identifier` where
/// the identifier names a recognized user-input source is D1.c.
fn dict_has_user_input_splat(node: Node<'_>, source: &[u8]) -> bool {
    if node.kind() != "dictionary" {
        return false;
    }
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if child.kind() != "dictionary_splat" {
            continue;
        }
        // `dictionary_splat` has the form `** <expression>`. Inspect
        // the expression text.
        let splat_text = node_text(child, source).unwrap_or("");
        // Strip leading `**` and whitespace.
        let expr = splat_text.trim_start_matches('*').trim();
        if is_user_input_identifier(expr) {
            return true;
        }
    }
    false
}

/// True iff the identifier text names a recognized user-input source
/// (matches one of the lexicon substrings).
fn is_user_input_identifier(expr: &str) -> bool {
    let lower = expr.to_lowercase();
    UNSTRUCTURED_JSON_USER_INPUT_SUBSTRINGS
        .iter()
        .any(|s| lower.contains(s))
}

/// True iff the dict literal contains a dangerous-operator key whose
/// value derives from user input.
///
/// Honest-review-driven heuristic: we don't fully trace data flow for
/// the operator value — we check whether the *dict text* contains
/// both a dangerous operator AND a user-input identifier in the same
/// dict literal. False positives are possible (the operator and the
/// user input could be in disjoint sub-keys) but the v0 conservative
/// stance is appropriate for a Critical-severity collapse.
fn has_dangerous_operator_with_user_input(_arg: Node<'_>, _source: &[u8], dict_text: &str) -> bool {
    if !line_contains_dangerous_operator(dict_text) {
        return false;
    }
    let lower = dict_text.to_lowercase();
    // Check whether any user-input identifier appears in the dict
    // text. We treat both TypedString and UnstructuredJson sources as
    // operator-injection-triggering when paired with $where, because
    // f-string interpolation of even a TypedString value into a $where
    // JavaScript expression is RCE.
    for s in UNSTRUCTURED_JSON_USER_INPUT_SUBSTRINGS {
        if lower.contains(s) {
            return true;
        }
    }
    for s in super::predict::TYPED_STRING_USER_INPUT_SUBSTRINGS {
        if lower.contains(s) {
            return true;
        }
    }
    false
}

/// True iff every user-input value position in the dict literal is
/// wrapped in a type-narrowing cast (str / ObjectId / int / float /
/// pydantic). If the dict has no user input, returns true (trivially
/// safe). Used by `classify_query_shape` to decide D1.a.
///
/// Implementation: pair-wise scan of the dict's `pair` children. For
/// each value node, check whether its text either (a) contains a
/// recognized type cast, or (b) does not contain any user-input
/// identifier (i.e., it's a literal or a non-user-derived variable).
fn is_typed_value_query(dict_node: Node<'_>, source: &[u8]) -> bool {
    let mut cursor = dict_node.walk();
    for child in dict_node.children(&mut cursor) {
        if child.kind() != "pair" {
            continue;
        }
        let Some(value) = child.child_by_field_name("value") else {
            continue;
        };
        let value_text = node_text(value, source).unwrap_or("");
        // If the value contains user-input identifier but does NOT
        // contain a type cast, the dict is not a typed-value query.
        if value_contains_user_input(value_text) && !value_contains_type_cast(value_text) {
            return false;
        }
        // Recurse into nested dict values (operator subdocs).
        if value.kind() == "dictionary" && !is_typed_value_query(value, source) {
            return false;
        }
    }
    true
}

fn value_contains_user_input(text: &str) -> bool {
    let lower = text.to_lowercase();
    for s in UNSTRUCTURED_JSON_USER_INPUT_SUBSTRINGS {
        if lower.contains(s) {
            return true;
        }
    }
    for s in super::predict::TYPED_STRING_USER_INPUT_SUBSTRINGS {
        if lower.contains(s) {
            return true;
        }
    }
    false
}

fn value_contains_type_cast(text: &str) -> bool {
    // Local, stricter list of value-wrap casts (including `str(`, which
    // is too broad as a free-line substring but is recognized when
    // wrapping an explicit dict-value expression in this evidence-
    // extraction context). Used only inside `is_typed_value_query`'s
    // value-position scan.
    const VALUE_CAST_SUBSTRINGS: &[&str] = &[
        "str(",
        "ObjectId(",
        "bson.ObjectId(",
        "int(",
        "float(",
        "bool(",
        "UUID(",
        "uuid.UUID(",
        ".parse_obj(",
        ".model_validate(",
        "schema.load(",
    ];
    VALUE_CAST_SUBSTRINGS.iter().any(|s| text.contains(s))
}

// ─────────────────────────────────────────────────────────────────────────────
// Extract full evidence from a classified call site
// ─────────────────────────────────────────────────────────────────────────────

/// Extract typed evidence from a Python pymongo query call node.
///
/// `call_node` must be a `call` AST node whose callee chain ends with
/// a pymongo query method name (post-`collect_python_nosql_sites`
/// filter). `module_root` is the file's module-level root node.
/// `source` is the file's raw bytes. `lines` is the pre-split
/// source-line slice the scanner already builds; used for the
/// annotation parsing and the user-input proximity check.
///
/// `file_path` is the source file path string; populated into the
/// returned `Evidence.file_path` for diagnostics. `api` is the
/// classification produced by `collect_python_nosql_sites`.
///
/// Never panics; missing fields produce defaults.
pub(super) fn extract_python_evidence<'a>(
    call_node: Node<'a>,
    _module_root: Node<'a>,
    source: &'a [u8],
    lines: &[&str],
    file_path: Option<String>,
    api: NosqlApi,
    callee_label: String,
) -> Evidence {
    let mut ev = Evidence {
        file_path,
        api: Some(api),
        callee_label: Some(callee_label),
        ..Default::default()
    };

    // ── Enclosing function (for name) and class. ──
    let fn_node = enclosing_python_function(call_node);
    if let Some(fn_node) = fn_node {
        if let Some(name_node) = fn_node.child_by_field_name("name") {
            if let Some(name) = node_text(name_node, source) {
                ev.enclosing_function = Some(name.to_string());
            }
        }
    }
    ev.enclosing_class = enclosing_python_class_name(call_node, source);

    // ── Function-name-based signals. ──
    if let Some(fn_name) = &ev.enclosing_function {
        ev.trust_boundary_name = matches_trust_boundary_name(fn_name);
        if matches_route_handler_name(fn_name) {
            ev.enclosing_route_handler = true;
        }
    }

    // ── Decorator-based route-handler detection. ──
    if !ev.enclosing_route_handler {
        if let Some(fn_node) = fn_node {
            if function_has_route_decorator(fn_node, lines) {
                ev.enclosing_route_handler = true;
            }
        }
    }

    // ── User-input source classification (±10 lines, prioritizing the
    //    call line itself). ──
    let call_line = call_node.start_position().row;
    ev.user_input_source = classify_nearby_user_input(lines, call_line, 10);

    // ── Type-cast proximity (±5 lines). ──
    ev.type_cast_nearby = type_cast_nearby(lines, call_line, 5);

    // ── Dollar-regex with user input. ──
    if let Some(line) = lines.get(call_line) {
        ev.has_dollar_regex_with_user_input =
            line.contains("$regex") && line_user_input_present(line);
    }

    // ── Developer-written operator (only if no user input in dict). ──
    // Inspect the call's first positional argument text for
    // developer-grade operators ($ne/$gt/$lt) with literal values.
    ev.has_developer_written_operator = has_developer_written_operator(call_node, source);

    // ── Source-line annotations. ──
    if let Some(line) = lines.get(call_line) {
        ev.nosql_safe_annotation = extract_nosql_safe_reason(line);
        ev.nosql_vulnerable_annotation = extract_nosql_vulnerable_source(line);
    }

    ev
}

/// True iff the line contains any recognized user-input identifier
/// (either source family).
fn line_user_input_present(line: &str) -> bool {
    !matches!(classify_user_input_source(line), UserInputSource::None)
}

/// Classify the user-input source family appearing in any line within
/// ±`radius` of `call_line` (0-indexed). Priority: UnstructuredJson
/// over TypedString (per the predict.rs `classify_user_input_source`).
fn classify_nearby_user_input(lines: &[&str], call_line: usize, radius: usize) -> UserInputSource {
    let start = call_line.saturating_sub(radius);
    let end = (call_line + radius + 1).min(lines.len());
    // First pass: prefer UnstructuredJson if it appears anywhere
    // in the window.
    for line in &lines[start..end] {
        if matches!(
            classify_user_input_source(line),
            UserInputSource::UnstructuredJson
        ) {
            return UserInputSource::UnstructuredJson;
        }
    }
    // Second pass: TypedString.
    for line in &lines[start..end] {
        if matches!(
            classify_user_input_source(line),
            UserInputSource::TypedString
        ) {
            return UserInputSource::TypedString;
        }
    }
    UserInputSource::None
}

/// True iff any line within ±`radius` of `call_line` contains a
/// recognized type-narrowing cast or pydantic validation.
fn type_cast_nearby(lines: &[&str], call_line: usize, radius: usize) -> bool {
    let start = call_line.saturating_sub(radius);
    let end = (call_line + radius + 1).min(lines.len());
    for line in &lines[start..end] {
        if line_contains_type_cast(line) {
            return true;
        }
    }
    false
}

/// True iff the call's first positional dict-literal argument contains
/// a developer-grade operator (`$ne`/`$gt`/`$lt`/`$in`) AND no user-
/// input identifier appears in the dict text. The "literal value"
/// condition is approximated by "no user-input identifier appears in
/// the dict literal text" — a conservative v0 heuristic.
fn has_developer_written_operator(call_node: Node<'_>, source: &[u8]) -> bool {
    let Some(args) = call_node.child_by_field_name("arguments") else {
        return false;
    };
    let Some(first_arg) = first_positional_arg(args) else {
        return false;
    };
    let dict_text = node_text(first_arg, source).unwrap_or("");
    if !line_contains_developer_operator(dict_text) {
        return false;
    }
    // No user input → literal value → developer-written.
    !value_contains_user_input(dict_text)
}

// ─────────────────────────────────────────────────────────────────────────────
// Decorator inspection
// ─────────────────────────────────────────────────────────────────────────────

/// True iff the `function_definition` node (Python) is preceded by a
/// `decorated_definition` whose decorator list contains a recognized
/// route-handler decorator. The check inspects the source-line text
/// because tree-sitter's grammar wraps decorated functions in a
/// `decorated_definition` node — we walk up one level to look.
fn function_has_route_decorator(fn_node: Node<'_>, lines: &[&str]) -> bool {
    let mut parent = fn_node.parent();
    while let Some(p) = parent {
        if p.kind() == "decorated_definition" {
            let mut cursor = p.walk();
            for child in p.children(&mut cursor) {
                if child.kind() == "decorator" {
                    let line_idx = child.start_position().row;
                    if let Some(line) = lines.get(line_idx) {
                        if matches_route_handler_decorator(line) {
                            return true;
                        }
                    }
                }
            }
            break;
        }
        if p.kind() == "module" {
            break;
        }
        parent = p.parent();
    }
    false
}

// ─────────────────────────────────────────────────────────────────────────────
// Enclosing class
// ─────────────────────────────────────────────────────────────────────────────

fn enclosing_python_class_name<'a>(node: Node<'a>, source: &'a [u8]) -> Option<String> {
    let mut cur = node.parent()?;
    loop {
        if cur.kind() == "class_definition" {
            let name = cur.child_by_field_name("name")?;
            return node_text(name, source).map(str::to_string);
        }
        if cur.kind() == "module" {
            return None;
        }
        cur = cur.parent()?;
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::detectors::ast_fingerprint::parse_root_ext;
    use crate::parsers::lightweight::Language;

    /// Parse `source` as Python and find the first `call` node whose
    /// function chain ends with the given attribute or identifier
    /// name.
    fn first_call_with_attr<'tree>(
        tree: &'tree tree_sitter::Tree,
        source: &[u8],
        attr_name: &str,
    ) -> tree_sitter::Node<'tree> {
        fn walk<'a>(
            node: tree_sitter::Node<'a>,
            source: &[u8],
            attr_name: &str,
        ) -> Option<tree_sitter::Node<'a>> {
            if node.kind() == "call" {
                if let Some(func) = node.child_by_field_name("function") {
                    let text = node_text(func, source).unwrap_or("");
                    let last = text.rsplit('.').next().unwrap_or("");
                    if last == attr_name {
                        return Some(node);
                    }
                }
            }
            let mut cursor = node.walk();
            for child in node.children(&mut cursor) {
                if let Some(found) = walk(child, source, attr_name) {
                    return Some(found);
                }
            }
            None
        }
        walk(tree.root_node(), source, attr_name)
            .unwrap_or_else(|| panic!("no call to {} found in source", attr_name))
    }

    fn extract(src: &str, attr: &str) -> Evidence {
        let tree = parse_root_ext(src, Language::Python, "py").expect("parse python");
        let root = tree.root_node();
        let call = first_call_with_attr(&tree, src.as_bytes(), attr);
        let lines: Vec<&str> = src.lines().collect();
        let func_text = call
            .child_by_field_name("function")
            .and_then(|f| node_text(f, src.as_bytes()))
            .unwrap_or("")
            .to_string();
        let api = classify_query_shape(call, src.as_bytes());
        extract_python_evidence(call, root, src.as_bytes(), &lines, None, api, func_text)
    }

    fn collect_sites(src: &str) -> Vec<(NosqlApi, String)> {
        let tree = parse_root_ext(src, Language::Python, "py").expect("parse python");
        let root = tree.root_node();
        collect_python_nosql_sites(root, src.as_bytes())
            .into_iter()
            .map(|s| (s.api, s.callee_label))
            .collect()
    }

    // ─── collect_python_nosql_sites: receiver-method shape ───
    #[test]
    fn collect_picks_up_find_one() {
        let sites = collect_sites("users.find_one({})\n");
        assert_eq!(sites.len(), 1);
        assert_eq!(sites[0].1, "users.find_one");
    }

    #[test]
    fn collect_picks_up_db_collection_aggregate() {
        let sites = collect_sites("db.users.aggregate([])\n");
        assert_eq!(sites.len(), 1);
        assert_eq!(sites[0].1, "db.users.aggregate");
    }

    #[test]
    fn collect_skips_bare_find_call() {
        // Bare `find(...)` without a receiver — not a method call.
        let sites = collect_sites("find()\n");
        assert!(sites.is_empty(), "bare find() must not be picked up");
    }

    #[test]
    fn collect_skips_items_find_array_method() {
        // Preserves the legacy `is_array_method` FP filter.
        let sites = collect_sites("items.find(predicate)\n");
        assert!(sites.is_empty(), "items.find must be skipped (array FP)");
    }

    #[test]
    fn collect_skips_list_find_array_method() {
        let sites = collect_sites("results.find(predicate)\n");
        assert!(sites.is_empty(), "results.find must be skipped");
    }

    // ─── Structural classification: D1.c DictExpansion ───
    #[test]
    fn classify_dict_expansion_request_json() {
        let sites = collect_sites("users.find_one({**request.json})\n");
        assert_eq!(sites[0].0, NosqlApi::DictExpansion);
    }

    #[test]
    fn classify_dict_expansion_request_body() {
        let sites = collect_sites("users.find({**request.body})\n");
        assert_eq!(sites[0].0, NosqlApi::DictExpansion);
    }

    #[test]
    fn classify_dict_expansion_get_json() {
        let sites = collect_sites("users.find_one({**request.get_json()})\n");
        assert_eq!(sites[0].0, NosqlApi::DictExpansion);
    }

    #[test]
    fn classify_dict_expansion_with_other_keys() {
        // **req.json mixed with other keys still triggers DictExpansion.
        let sites = collect_sites("users.find_one({**request.json, \"active\": True})\n");
        assert_eq!(sites[0].0, NosqlApi::DictExpansion);
    }

    // ─── Structural classification: D1.b OperatorInjection ───
    #[test]
    fn classify_where_with_user_input() {
        let src = "users.find_one({\"$where\": f\"this.x == '{request.form['x']}'\"})\n";
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::OperatorInjection);
    }

    #[test]
    fn classify_function_operator_with_user_input() {
        let src = "users.aggregate([{\"$function\": {\"body\": request.json['code']}}])\n";
        // Aggregate-pipeline first arg is a list, not a dict. The
        // `has_dangerous_operator_with_user_input` text-level check
        // still fires because the list's source text contains both
        // `$function` and `request.json` — substring co-occurrence
        // suffices for D1.b in v0. This is conservative (a list
        // wrapping a dangerous-operator pipeline stage IS the
        // OperatorInjection collapse).
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::OperatorInjection);
    }

    #[test]
    fn classify_where_without_user_input_is_ambiguous() {
        // $where with no user input near it — still suspicious in
        // real code, but the v0 OperatorInjection collapse requires
        // user input. Falls back to Ambiguous (weighted scoring).
        let src = "users.find({\"$where\": \"this.x == 1\"})\n";
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::Ambiguous);
    }

    // ─── Structural classification: D1.a TypedValueQuery ───
    #[test]
    fn classify_typed_value_query_with_str_cast() {
        let src = "users.find_one({\"username\": str(request.form['user'])})\n";
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::TypedValueQuery);
    }

    #[test]
    fn classify_typed_value_query_with_objectid_cast() {
        let src = "users.find({\"_id\": ObjectId(request.form['id'])})\n";
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::TypedValueQuery);
    }

    #[test]
    fn classify_typed_value_query_with_no_user_input() {
        let src = "users.find({\"active\": True, \"role\": \"admin\"})\n";
        let sites = collect_sites(src);
        // No user input → trivially typed-value safe.
        assert_eq!(sites[0].0, NosqlApi::TypedValueQuery);
    }

    #[test]
    fn classify_typed_value_query_with_developer_operator() {
        // `$ne` with literal value — typed value query (no dangerous
        // operators, no user input in dict).
        let src = "users.find({\"role\": {\"$ne\": \"admin\"}})\n";
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::TypedValueQuery);
    }

    // ─── Structural classification: Ambiguous fallback ───
    #[test]
    fn classify_naked_request_json_value_is_ambiguous() {
        // The D5.2 honest-review Case A: user_json in value position,
        // no cast. Falls to Ambiguous (weighted scoring will tip
        // RealBug via UnstructuredJson source + handler).
        let src = "users.find_one({\"username\": request.json['user']})\n";
        let sites = collect_sites(src);
        assert_eq!(sites[0].0, NosqlApi::Ambiguous);
    }

    #[test]
    fn classify_identifier_arg_is_ambiguous() {
        // First arg is an identifier — predictor can't trace; Ambiguous.
        let src = "q = build_query()\nusers.find(q)\n";
        let sites = collect_sites(src);
        // collect_sites is a flat scanner; both calls (build_query
        // isn't pymongo) — only users.find is collected.
        let pymongo: Vec<_> = sites.iter().filter(|(_, l)| l == "users.find").collect();
        assert_eq!(pymongo.len(), 1);
        assert_eq!(pymongo[0].0, NosqlApi::Ambiguous);
    }

    // ─── Enclosing function detection ───
    #[test]
    fn detects_enclosing_function_name() {
        let src = "\
            def get_user(user_id):\n\
            \x20   return users.find_one({\"_id\": ObjectId(user_id)})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.enclosing_function.as_deref(), Some("get_user"));
    }

    #[test]
    fn detects_enclosing_class_name() {
        let src = "\
            class UserRepo:\n\
            \x20   def get(self, uid):\n\
            \x20       return users.find_one({\"_id\": ObjectId(uid)})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.enclosing_class.as_deref(), Some("UserRepo"));
    }

    // ─── Decorator-based route-handler detection ───
    #[test]
    fn detects_flask_route_decorator() {
        let src = "\
            from flask import request\n\
            @app.route('/login', methods=['POST'])\n\
            def login():\n\
            \x20   return users.find_one({\"name\": request.form['n']})\n";
        let ev = extract(src, "find_one");
        assert!(ev.enclosing_route_handler);
    }

    #[test]
    fn detects_handler_name() {
        let src = "\
            def login_handler():\n\
            \x20   return users.find_one({})\n";
        let ev = extract(src, "find_one");
        assert!(ev.enclosing_route_handler);
    }

    // ─── User-input source classification ───
    #[test]
    fn detects_unstructured_json_source() {
        let src = "\
            def f():\n\
            \x20   body = request.json\n\
            \x20   return users.find_one({\"x\": body['x']})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.user_input_source, UserInputSource::UnstructuredJson);
    }

    #[test]
    fn detects_typed_string_source() {
        let src = "\
            def f():\n\
            \x20   n = request.form['name']\n\
            \x20   return users.find_one({\"name\": str(n)})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.user_input_source, UserInputSource::TypedString);
    }

    #[test]
    fn detects_no_user_input_source() {
        let src = "\
            def f():\n\
            \x20   return users.find_one({\"active\": True})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.user_input_source, UserInputSource::None);
    }

    // ─── Type-cast nearby detection ───
    #[test]
    fn detects_objectid_cast_nearby() {
        let src = "\
            def f(uid):\n\
            \x20   oid = ObjectId(uid)\n\
            \x20   return users.find_one({\"_id\": oid})\n";
        let ev = extract(src, "find_one");
        assert!(ev.type_cast_nearby);
    }

    #[test]
    fn detects_pydantic_validation_nearby() {
        let src = "\
            def f(payload):\n\
            \x20   q = QuerySchema.model_validate(payload)\n\
            \x20   return users.find_one({\"name\": q.name})\n";
        let ev = extract(src, "find_one");
        assert!(ev.type_cast_nearby);
    }

    // ─── Developer-written operator detection ───
    #[test]
    fn detects_developer_written_ne() {
        // Case D headline: `$ne` with literal value, no user input
        // → developer-written.
        let src = "users.find({\"role\": {\"$ne\": \"admin\"}})\n";
        let ev = extract(src, "find");
        assert!(ev.has_developer_written_operator);
    }

    #[test]
    fn does_not_detect_developer_written_when_user_input_present() {
        // `$ne` with user input — NOT developer-written; this is a
        // legitimate operator-injection concern (though not Critical-
        // grade like $where).
        let src = "users.find({\"role\": {\"$ne\": request.form['r']}})\n";
        let ev = extract(src, "find");
        assert!(!ev.has_developer_written_operator);
    }

    #[test]
    fn does_not_detect_developer_written_for_plain_query() {
        // No operator at all — no developer-written-operator signal.
        let src = "users.find({\"role\": \"admin\"})\n";
        let ev = extract(src, "find");
        assert!(!ev.has_developer_written_operator);
    }

    // ─── Annotation parsing ───
    #[test]
    fn extracts_nosql_safe_annotation() {
        let src = "\
            def f():\n\
            \x20   return users.find_one(q)  # repotoire: nosql-safe[pydantic-validated]\n";
        let ev = extract(src, "find_one");
        assert_eq!(
            ev.nosql_safe_annotation.as_deref(),
            Some("pydantic-validated")
        );
    }

    #[test]
    fn extracts_nosql_vulnerable_annotation() {
        let src = "\
            def f():\n\
            \x20   return users.find_one(q)  # repotoire: nosql-vulnerable[helper-built]\n";
        let ev = extract(src, "find_one");
        assert_eq!(
            ev.nosql_vulnerable_annotation.as_deref(),
            Some("helper-built")
        );
    }

    #[test]
    fn does_not_extract_unrelated_annotation() {
        let src = "\
            def f():\n\
            \x20   return users.find_one({})  # repotoire: jwt-safe[verified]\n";
        let ev = extract(src, "find_one");
        assert!(ev.nosql_safe_annotation.is_none());
        assert!(ev.nosql_vulnerable_annotation.is_none());
    }

    // ─── End-to-end integration: full evidence for §6 Case A ───
    #[test]
    fn case_a_full_evidence_naked_request_json() {
        let src = "\
            from flask import request\n\
            @app.route('/login', methods=['POST'])\n\
            def login():\n\
            \x20   return users.find_one({\"username\": request.json['user']})\n";
        let ev = extract(src, "find_one");
        // Naked user input in value position — Ambiguous classification,
        // weighted scoring will tip RealBug.
        assert_eq!(ev.api, Some(NosqlApi::Ambiguous));
        assert_eq!(ev.user_input_source, UserInputSource::UnstructuredJson);
        assert!(ev.enclosing_route_handler);
    }

    // ─── End-to-end integration: full evidence for §6 Case B ───
    #[test]
    fn case_b_full_evidence_str_cast() {
        let src = "\
            from flask import request\n\
            def login():\n\
            \x20   return users.find_one({\"username\": str(request.form['user'])})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.api, Some(NosqlApi::TypedValueQuery));
        // The user input source still gets classified — TypedString
        // (request.form), but the API collapse dominates the scoring.
        assert_eq!(ev.user_input_source, UserInputSource::TypedString);
    }

    // ─── End-to-end integration: full evidence for §6 Case C ───
    #[test]
    fn case_c_full_evidence_where_with_user_input() {
        let src = "\
            from flask import request\n\
            def f():\n\
            \x20   return users.find_one({\"$where\": f\"this.x == '{request.form['x']}'\"})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.api, Some(NosqlApi::OperatorInjection));
    }

    // ─── End-to-end integration: full evidence for §6 Case D (FP REDUCTION) ───
    #[test]
    fn case_d_full_evidence_developer_written_ne() {
        let src = "\
            def list_users():\n\
            \x20   return users.find({\"role\": {\"$ne\": \"admin\"}})\n";
        let ev = extract(src, "find");
        // TypedValueQuery classification — no user input, no
        // dangerous operators → D1.a collapse fires.
        assert_eq!(ev.api, Some(NosqlApi::TypedValueQuery));
        assert!(ev.has_developer_written_operator);
    }

    // ─── End-to-end integration: full evidence for §6 Case E ───
    #[test]
    fn case_e_full_evidence_dict_expansion() {
        let src = "\
            from flask import request\n\
            @app.route('/q', methods=['POST'])\n\
            def f():\n\
            \x20   return users.find_one({**request.get_json()})\n";
        let ev = extract(src, "find_one");
        assert_eq!(ev.api, Some(NosqlApi::DictExpansion));
        assert!(ev.enclosing_route_handler);
    }

    // ─── End-to-end integration: full evidence for §6 Case F ───
    #[test]
    fn case_f_full_evidence_objectid_cast() {
        let src = "\
            from flask import request\n\
            def f():\n\
            \x20   return users.find({\"_id\": ObjectId(request.form['id'])})\n";
        let ev = extract(src, "find");
        assert_eq!(ev.api, Some(NosqlApi::TypedValueQuery));
        assert!(ev.type_cast_nearby);
    }
}