repotoire 0.9.0

//! AST-driven extraction of [`super::predict::Evidence`] for Python
//! weak-hash 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.
//!
//! Splitting the two halves matches the typical Phase 1 ↔ Phase 2 split
//! in this codebase: the predictor knows about weights and signals;
//! the extractor knows about Python AST shapes. Either can be revised
//! without touching the other.
//!
//! # What this module knows about
//!
//! - Walking up from the call node to the enclosing
//!   `function_definition` and `class_definition`.
//! - Reading the first positional argument and deciding whether it's
//!   a bare identifier (vs a string literal, expression, call, ...).
//! - Detecting `usedforsecurity=False` as a literal boolean kwarg.
//! - Detecting `[:N]` truncation on the call result by looking at the
//!   call's parent / grandparent for a `subscript` with a `slice`.
//! - Detecting `os.urandom(...)` anywhere inside the first positional
//!   argument expression (so `hashlib.sha1(os.urandom(8) + nonce)`
//!   counts).
//! - Reading the source line for a trailing
//!   `# repotoire: protocol-required[<rfc>]` annotation.
//!
//! # What this module deliberately does NOT do
//!
//! - Does not look for evidence in non-Python languages. Phase 2a
//!   scope is Python-only per the decisions doc.
//! - Does not try to follow data flow to figure out what `s` actually
//!   is. The first-arg-identifier signal is "what's the bare name."
//!   Anything fancier is a future Phase.
//! - Does not consult the graph (Phase 1c) for enclosing scope. AST
//!   walking up to `class_definition` is sufficient and avoids a
//!   detector → graph dependency we don't need yet.

use super::predict::{extract_protocol_required_rfc, Evidence};
use crate::detectors::security::ast_helpers::{
    collect_named_args, enclosing_python_function, node_text,
};
use tree_sitter::Node;

/// Extract typed evidence from a Python call node.
///
/// `call_node` must be a `call` AST node whose function names a weak
/// hash primitive. `source` is the file's raw bytes. `lines` is the
/// pre-split source-line slice the scanner already builds; we use it
/// only to read the source line for annotation lookup.
///
/// The function never panics; missing fields produce `None`/`false`
/// in the corresponding Evidence field.
pub(super) fn extract_python_evidence<'a>(
    call_node: Node<'a>,
    source: &'a [u8],
    lines: &[&str],
) -> Evidence {
    let mut ev = Evidence::default();

    // ── Enclosing function and class. ──
    if let Some(fn_node) = enclosing_python_function(call_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);

    // ── First positional argument. ──
    if let Some(args_node) = call_node.child_by_field_name("arguments") {
        let args = collect_named_args(args_node);
        let first_positional = args
            .iter()
            .find(|a| a.kind() != "keyword_argument" && a.kind() != "comment")
            .copied();
        if let Some(arg) = first_positional {
            ev.first_arg_ident = python_arg_as_identifier(arg, source);
            ev.input_includes_urandom = expression_contains_os_urandom(arg, source);
        }

        // ── usedforsecurity=False kwarg (collapsing). ──
        ev.usedforsecurity_false = has_usedforsecurity_false(&args, source);
    }

    // ── Result truncation [:N]. ──
    ev.result_truncated = call_result_is_sliced(call_node);

    // ── Source-line annotation. ──
    let line_idx = call_node.start_position().row;
    if let Some(line) = lines.get(line_idx) {
        ev.protocol_required_annotation = extract_protocol_required_rfc(line);
    }

    ev
}

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

/// Walk up from `node` to the nearest enclosing `class_definition` and
/// return its name. Returns `None` at module level or when only enclosed
/// by functions.
///
/// Stops at `module` (top of file) without going further. Methods inside
/// a class return that class's name even when nested in a function (e.g.
/// inner closures); this matches the user's mental model of "what class
/// am I in."
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()?;
    }
}

/// If `arg` is a bare identifier expression, return its name. Returns
/// `None` for string literals, calls, binary expressions, attribute
/// accesses, parenthesized exprs, and anything else.
///
/// We deliberately do NOT unwrap `(x)` to `x` here — a parenthesized
/// expression is a code smell in this context, and the predictor's
/// "non-sensitive bonus" should not fire on something the author
/// went out of their way to wrap. If misprediction logging shows
/// false RealBugs from this, revisit.
fn python_arg_as_identifier<'a>(arg: Node<'a>, source: &'a [u8]) -> Option<String> {
    if arg.kind() == "identifier" {
        node_text(arg, source).map(str::to_string)
    } else {
        None
    }
}

/// True if `expr` (or any of its descendants) is a call to
/// `os.urandom(...)`.
///
/// Walks the entire subtree; matching is by attribute name + receiver.
/// This catches `os.urandom(8)`, `os.urandom(8) + nonce`,
/// `b64encode(os.urandom(8))`, etc.
fn expression_contains_os_urandom<'a>(expr: Node<'a>, source: &'a [u8]) -> bool {
    if is_os_urandom_call(expr, source) {
        return true;
    }
    let mut cursor = expr.walk();
    for child in expr.children(&mut cursor) {
        if expression_contains_os_urandom(child, source) {
            return true;
        }
    }
    false
}

fn is_os_urandom_call<'a>(node: Node<'a>, source: &'a [u8]) -> bool {
    if node.kind() != "call" {
        return false;
    }
    let Some(func) = node.child_by_field_name("function") else {
        return false;
    };
    if func.kind() != "attribute" {
        return false;
    }
    let Some(obj) = func.child_by_field_name("object") else {
        return false;
    };
    let Some(attr) = func.child_by_field_name("attribute") else {
        return false;
    };
    node_text(obj, source) == Some("os") && node_text(attr, source) == Some("urandom")
}

/// True if `args` contains a `usedforsecurity=False` keyword argument
/// where the value is the literal `False`.
///
/// A non-literal value (e.g. `usedforsecurity=some_var`) does NOT
/// count — we don't try to do data-flow.
fn has_usedforsecurity_false(args: &[Node<'_>], source: &[u8]) -> bool {
    use crate::detectors::security::ast_helpers::python_kwarg_value;
    let Some(value) = python_kwarg_value(args, "usedforsecurity", source) else {
        return false;
    };
    // Tree-sitter Python uses `false` as the kind for the `False` keyword.
    value.kind() == "false"
}

/// True if the call's result is sliced with `[:N]` somewhere in its
/// surrounding expression.
///
/// Walks UP from the call node looking for a `subscript` whose
/// `subscript`/`slice` child matches `[:N]` shape. Stops at
/// statement boundaries. Handles common patterns:
///
/// - `hashlib.sha1(s).hexdigest()[:8]`
/// - `hashlib.sha1(s).digest()[:16]`
/// - `hashlib.sha1(s)[:8]` (less common but valid AST)
///
/// Does NOT recognize `.hexdigest()[0:N]` differently from `[:N]`;
/// either truncation pattern counts.
fn call_result_is_sliced(call_node: Node<'_>) -> bool {
    let mut cur = call_node;
    // Walk up at most a few levels — past `.hexdigest()` etc — until
    // we either find a subscript or leave the expression.
    for _ in 0..5 {
        let Some(parent) = cur.parent() else {
            return false;
        };
        match parent.kind() {
            "subscript" => {
                // The subscript node has a `subscript` field that is
                // either an int, a slice, or other. We only care about
                // slice children.
                let mut cursor = parent.walk();
                for child in parent.children(&mut cursor) {
                    if child.kind() == "slice" {
                        return true;
                    }
                }
                return false;
            }
            // Wrapper nodes that don't end the expression.
            "attribute" | "call" | "parenthesized_expression" => {
                cur = parent;
            }
            // We've reached a statement-level node: no slice.
            _ => return false,
        }
    }
    false
}

// ─────────────────────────────────────────────────────────────────────────────
// 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 in the named attribute (e.g. `"sha1"` for
    /// `hashlib.sha1(...)`). Used to anchor evidence-extraction tests
    /// on a specific call site without index-counting.
    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") {
                    if func.kind() == "attribute" {
                        if let Some(attr) = func.child_by_field_name("attribute") {
                            if node_text(attr, source) == Some(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 ending in `.{attr_name}` in source"))
    }

    fn extract(source: &str, attr: &str) -> Evidence {
        let bytes = source.as_bytes();
        let tree = parse_root_ext(source, Language::Python, "py").expect("parse python");
        let lines: Vec<&str> = source.lines().collect();
        let call = first_call_with_attr(&tree, bytes, attr);
        extract_python_evidence(call, bytes, &lines)
    }

    // ── Enclosing scope ──

    #[test]
    fn extracts_enclosing_class_and_function() {
        let src = "
class DigestAuth:
    def _get_client_nonce(self, s):
        return hashlib.sha1(s).hexdigest()[:8]
";
        let ev = extract(src, "sha1");
        assert_eq!(ev.enclosing_class.as_deref(), Some("DigestAuth"));
        assert_eq!(ev.enclosing_function.as_deref(), Some("_get_client_nonce"));
    }

    #[test]
    fn module_level_call_has_no_enclosing_class() {
        let src = "import hashlib\nh = hashlib.sha1(b'data')\n";
        let ev = extract(src, "sha1");
        assert!(ev.enclosing_class.is_none());
        assert!(ev.enclosing_function.is_none());
    }

    #[test]
    fn function_level_call_has_no_enclosing_class() {
        let src = "
def hash_thing(x):
    return hashlib.sha1(x).digest()
";
        let ev = extract(src, "sha1");
        assert!(ev.enclosing_class.is_none());
        assert_eq!(ev.enclosing_function.as_deref(), Some("hash_thing"));
    }

    // ── First positional argument ──

    #[test]
    fn first_arg_identifier_extracted() {
        let src = "import hashlib\nh = hashlib.sha1(password)\n";
        let ev = extract(src, "sha1");
        assert_eq!(ev.first_arg_ident.as_deref(), Some("password"));
    }

    #[test]
    fn first_arg_string_literal_yields_no_identifier() {
        let src = "import hashlib\nh = hashlib.sha1(b'data')\n";
        let ev = extract(src, "sha1");
        assert!(ev.first_arg_ident.is_none());
    }

    #[test]
    fn first_arg_expression_yields_no_identifier() {
        let src = "import hashlib\nh = hashlib.sha1(s + nonce)\n";
        let ev = extract(src, "sha1");
        assert!(ev.first_arg_ident.is_none());
    }

    #[test]
    fn first_arg_call_yields_no_identifier() {
        let src = "import hashlib\nh = hashlib.sha1(get_password())\n";
        let ev = extract(src, "sha1");
        assert!(ev.first_arg_ident.is_none());
    }

    // ── Truncation ──

    #[test]
    fn hexdigest_truncated_is_detected() {
        let src = "import hashlib\nh = hashlib.sha1(s).hexdigest()[:8]\n";
        let ev = extract(src, "sha1");
        assert!(ev.result_truncated);
    }

    #[test]
    fn untruncated_hexdigest_is_not_detected() {
        let src = "import hashlib\nh = hashlib.sha1(s).hexdigest()\n";
        let ev = extract(src, "sha1");
        assert!(!ev.result_truncated);
    }

    #[test]
    fn digest_indexed_with_int_is_not_truncation() {
        // `[0]` is integer subscript, not slice. Single-byte access
        // is not "truncation"; it's something else (and probably a bug).
        let src = "import hashlib\nh = hashlib.sha1(s).digest()[0]\n";
        let ev = extract(src, "sha1");
        assert!(!ev.result_truncated);
    }

    // ── os.urandom in input ──

    #[test]
    fn input_with_os_urandom_is_detected() {
        let src = "import os, hashlib\nh = hashlib.sha1(os.urandom(8)).hexdigest()\n";
        let ev = extract(src, "sha1");
        assert!(ev.input_includes_urandom);
    }

    #[test]
    fn input_with_nested_os_urandom_is_detected() {
        let src = "import os, hashlib\nh = hashlib.sha1(os.urandom(8) + nonce).hexdigest()\n";
        let ev = extract(src, "sha1");
        assert!(ev.input_includes_urandom);
    }

    #[test]
    fn input_without_os_urandom_is_not_detected() {
        let src = "import hashlib\nh = hashlib.sha1(b'data').hexdigest()\n";
        let ev = extract(src, "sha1");
        assert!(!ev.input_includes_urandom);
    }

    // ── usedforsecurity=False ──

    #[test]
    fn usedforsecurity_false_kwarg_is_detected() {
        let src = "import hashlib\nh = hashlib.md5(b'data', usedforsecurity=False)\n";
        let ev = extract(src, "md5");
        assert!(ev.usedforsecurity_false);
    }

    #[test]
    fn usedforsecurity_true_is_not_detected_as_false() {
        let src = "import hashlib\nh = hashlib.md5(b'data', usedforsecurity=True)\n";
        let ev = extract(src, "md5");
        assert!(!ev.usedforsecurity_false);
    }

    #[test]
    fn usedforsecurity_variable_is_not_detected_as_false() {
        // `usedforsecurity=some_var` is a non-literal; we don't do
        // data flow. The kwarg must be the literal False.
        let src = "import hashlib\nh = hashlib.md5(b'data', usedforsecurity=flag)\n";
        let ev = extract(src, "md5");
        assert!(!ev.usedforsecurity_false);
    }

    #[test]
    fn missing_usedforsecurity_kwarg_is_false() {
        let src = "import hashlib\nh = hashlib.md5(b'data')\n";
        let ev = extract(src, "md5");
        assert!(!ev.usedforsecurity_false);
    }

    // ── Annotation ──

    #[test]
    fn protocol_required_annotation_extracted() {
        let src = "import hashlib\nh = hashlib.sha1(s)  # repotoire: protocol-required[RFC7616]\n";
        let ev = extract(src, "sha1");
        assert_eq!(ev.protocol_required_annotation.as_deref(), Some("RFC7616"),);
    }

    #[test]
    fn no_annotation_yields_none() {
        let src = "import hashlib\nh = hashlib.sha1(s)\n";
        let ev = extract(src, "sha1");
        assert!(ev.protocol_required_annotation.is_none());
    }

    #[test]
    fn unrelated_annotation_yields_none() {
        let src = "import hashlib\nh = hashlib.sha1(s)  # repotoire: low-entropy[md5]\n";
        let ev = extract(src, "sha1");
        assert!(ev.protocol_required_annotation.is_none());
    }

    // ── Integration: the worked example ──

    #[test]
    fn worked_example_extracts_all_four_signals() {
        // The decisions doc's worked example: HTTPX `_auth.py:309`
        // shape — DigestAuth class, _get_client_nonce method, first
        // arg `s` (non-sensitive identifier), result truncated, input
        // includes os.urandom.
        let src = "
import os
import hashlib

class DigestAuth:
    def _get_client_nonce(self):
        s = b''
        s = os.urandom(8) + s
        return hashlib.sha1(s).hexdigest()[:8]
";
        let ev = extract(src, "sha1");
        assert_eq!(ev.enclosing_class.as_deref(), Some("DigestAuth"));
        assert_eq!(ev.first_arg_ident.as_deref(), Some("s"));
        assert!(ev.result_truncated);
        // Note: input_includes_urandom checks the FIRST ARG expression,
        // not the surrounding function. `s` is just an identifier; its
        // assignment elsewhere uses os.urandom. The signal as designed
        // matches when urandom is INSIDE the call's first arg, not when
        // the variable was previously assigned from urandom. This is
        // the conservative interpretation; see decisions doc for why
        // we don't do data flow here.
        assert!(
            !ev.input_includes_urandom,
            "first_arg_ident=`s` doesn't itself contain os.urandom; data flow is out of scope"
        );
        assert!(!ev.usedforsecurity_false);
        assert!(ev.protocol_required_annotation.is_none());
    }

    #[test]
    fn worked_example_inline_extracts_all_four_signals() {
        // Same shape but with os.urandom() inlined as the call's
        // first arg. THIS is what the predictor's W_INPUT_INCLUDES_URANDOM
        // weight is calibrated for (no data flow needed).
        let src = "
import os
import hashlib

class DigestAuth:
    def _get_client_nonce(self):
        return hashlib.sha1(os.urandom(8)).hexdigest()[:8]
";
        let ev = extract(src, "sha1");
        assert_eq!(ev.enclosing_class.as_deref(), Some("DigestAuth"));
        // First arg is `os.urandom(8)`, not a bare identifier.
        assert!(ev.first_arg_ident.is_none());
        assert!(ev.result_truncated);
        assert!(ev.input_includes_urandom);
    }
}