pyrograph 0.1.0

//! Core unit tests for the taint analysis engine.
//! Tests graph construction, label loading, scale, and adversarial patterns.

use pyrograph::ir::{EdgeKind, NodeKind, TaintGraph};
use pyrograph::labels::{load_labels, label_node, TaintLabel};
use pyrograph::analyze;
use std::path::Path;

#[test]
fn scale_1000_node_chain() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let mut prev = src;
    for i in 1..=998 {
        let node = graph.add_node(NodeKind::Variable, format!("n{}", i).into(), None);
        graph.add_edge(prev, node, EdgeKind::Assignment);
        prev = node;
    }
    let sink = graph.add_node(NodeKind::Call, "sink".into(), Some(TaintLabel::Sink(0)));
    graph.add_edge(prev, sink, EdgeKind::Argument);
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1, "expected exactly 1 finding");
    assert_eq!(findings[0].path.len(), 1000, "expected path length 1000");
}

#[test]
fn scale_wide_fanout_100_sinks() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    for i in 0..100 {
        let sink = graph.add_node(NodeKind::Call, format!("sink{}", i).into(), Some(TaintLabel::Sink(i)));
        graph.add_edge(src, sink, EdgeKind::Argument);
    }
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 100, "expected exactly 100 findings");
}

#[test]
fn scale_100_sources_100_sinks_grid() {
    let mut graph = TaintGraph::new();
    let mut sources = Vec::new();
    for i in 0..100 {
        let src = graph.add_node(NodeKind::Variable, format!("src{}", i).into(), Some(TaintLabel::Source(i)));
        sources.push(src);
    }
    let mut sinks = Vec::new();
    for j in 0..100 {
        let sink = graph.add_node(NodeKind::Call, format!("sink{}", j).into(), Some(TaintLabel::Sink(j)));
        sinks.push(sink);
    }
    for src in &sources {
        for sink in &sinks {
            let mid = graph.add_node(NodeKind::Variable, "mid".into(), None);
            graph.add_edge(*src, mid, EdgeKind::Assignment);
            graph.add_edge(mid, *sink, EdgeKind::Argument);
        }
    }
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 10000, "expected exactly 10000 findings");
}

#[test]
fn adversarial_source_is_also_sink() {
    let mut graph = TaintGraph::new();
    let both = graph.add_node(NodeKind::Variable, "both".into(), Some(TaintLabel::Both(0, 0)));
    // No edges
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0, "Both without argument flow must not self-report");
    let _ = both;
}

#[test]
fn adversarial_deep_cycle_no_infinite_loop() {
    let mut graph = TaintGraph::new();
    let mut nodes = Vec::new();
    for i in 0..500 {
        let label = if i == 0 {
            Some(TaintLabel::Source(0))
        } else {
            None
        };
        let node = graph.add_node(NodeKind::Variable, format!("n{}", i).into(), label);
        nodes.push(node);
    }
    for i in 0..500 {
        let next = (i + 1) % 500;
        graph.add_edge(nodes[i], nodes[next], EdgeKind::Assignment);
    }
    let sink = graph.add_node(NodeKind::Call, "sink".into(), Some(TaintLabel::Sink(0)));
    graph.add_edge(nodes[250], sink, EdgeKind::Argument);
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1, "expected exactly 1 finding in cycle");
    assert_eq!(findings[0].path.len(), 252, "expected shortest path length 252 from node 0 to sink via node 250");
}

#[test]
fn adversarial_duplicate_edges() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let sink = graph.add_node(NodeKind::Call, "sink".into(), Some(TaintLabel::Sink(0)));
    for _ in 0..10 {
        graph.add_edge(src, sink, EdgeKind::Argument);
    }
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1, "duplicate edges must not create duplicate findings");
}

#[test]
fn adversarial_zero_node_graph() {
    let graph = TaintGraph::new();
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0, "empty graph must yield zero findings");
}

#[test]
fn adversarial_sink_before_source_in_node_order() {
    let mut graph = TaintGraph::new();
    let sink = graph.add_node(NodeKind::Call, "sink".into(), Some(TaintLabel::Sink(0)));
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    graph.add_edge(src, sink, EdgeKind::Argument);
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1, "must find flow regardless of node insertion order");
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
    assert_eq!(findings[0].path, vec![src, sink]);
}

#[test]
fn adversarial_all_nodes_are_sources() {
    let mut graph = TaintGraph::new();
    let mut sources = Vec::new();
    for i in 0..100 {
        let src = graph.add_node(NodeKind::Variable, format!("src{}", i).into(), Some(TaintLabel::Source(i)));
        sources.push(src);
    }
    let sink = graph.add_node(NodeKind::Call, "sink".into(), Some(TaintLabel::Sink(0)));
    for src in &sources {
        graph.add_edge(*src, sink, EdgeKind::Argument);
    }
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 100, "100 sources to 1 sink should yield 100 findings");
}

#[test]
fn adversarial_all_nodes_are_sinks() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let mut prev = src;
    for i in 0..100 {
        let sink = graph.add_node(NodeKind::Call, format!("sink{}", i).into(), Some(TaintLabel::Sink(i)));
        graph.add_edge(prev, sink, EdgeKind::Argument);
        prev = sink;
    }
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 100, "chain of 100 sinks from 1 source should yield 100 findings");
}

#[test]
fn basic_single_hop_source_to_sink() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "process.env".into(), Some(TaintLabel::Source(0)));
    let snk = graph.add_node(NodeKind::Call, "fetch".into(), Some(TaintLabel::Sink(0)));
    graph.add_edge(src, snk, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
    assert_eq!(findings[0].path, vec![src, snk]);
}

#[test]
fn basic_multi_hop_chain() {
    let mut graph = TaintGraph::new();
    let n1 = graph.add_node(NodeKind::Variable, "process.env".into(), Some(TaintLabel::Source(0)));
    let n2 = graph.add_node(NodeKind::Variable, "a".into(), None);
    let n3 = graph.add_node(NodeKind::Variable, "b".into(), None);
    let n4 = graph.add_node(NodeKind::Variable, "c".into(), None);
    let n5 = graph.add_node(NodeKind::Call, "eval".into(), Some(TaintLabel::Sink(0)));

    graph.add_edge(n1, n2, EdgeKind::Assignment);
    graph.add_edge(n2, n3, EdgeKind::Assignment);
    graph.add_edge(n3, n4, EdgeKind::Assignment);
    graph.add_edge(n4, n5, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
    assert_eq!(findings[0].path.len(), 5);
}

#[test]
fn basic_no_path_no_finding() {
    let mut graph = TaintGraph::new();
    let _src = graph.add_node(NodeKind::Variable, "process.env".into(), Some(TaintLabel::Source(0)));
    let _mid = graph.add_node(NodeKind::Variable, "x".into(), None);
    let _snk = graph.add_node(NodeKind::Call, "fetch".into(), Some(TaintLabel::Sink(0)));
    // No edges at all

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0);
}

#[test]
fn basic_disconnected_graph() {
    let mut graph = TaintGraph::new();
    // Component A: source -> intermediate -> sink (connected)
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let mid = graph.add_node(NodeKind::Variable, "mid".into(), None);
    let snk = graph.add_node(NodeKind::Call, "snk".into(), Some(TaintLabel::Sink(0)));
    graph.add_edge(src, mid, EdgeKind::Assignment);
    graph.add_edge(mid, snk, EdgeKind::Argument);

    // Component B: source and sink completely disconnected (no edges at all)
    let _src2 = graph.add_node(NodeKind::Variable, "src2".into(), Some(TaintLabel::Source(1)));
    let _orphan = graph.add_node(NodeKind::Variable, "orphan".into(), None);
    let _snk2 = graph.add_node(NodeKind::Call, "snk2".into(), Some(TaintLabel::Sink(1)));
    // Intentionally no edges involving src2 or snk2

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
}

#[test]
fn multiple_two_sources_one_sink() {
    let mut graph = TaintGraph::new();
    let s1 = graph.add_node(NodeKind::Variable, "s1".into(), Some(TaintLabel::Source(0)));
    let s2 = graph.add_node(NodeKind::Variable, "s2".into(), Some(TaintLabel::Source(1)));
    let snk = graph.add_node(NodeKind::Call, "fetch".into(), Some(TaintLabel::Sink(0)));

    graph.add_edge(s1, snk, EdgeKind::Argument);
    graph.add_edge(s2, snk, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 2);

    let sources: Vec<_> = findings.iter().map(|f| f.source).collect();
    let sinks: Vec<_> = findings.iter().map(|f| f.sink).collect();
    assert!(sources.contains(&0));
    assert!(sources.contains(&1));
    assert_eq!(sinks, vec![0, 0]);
}

#[test]
fn multiple_one_source_two_sinks() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let k1 = graph.add_node(NodeKind::Call, "fetch".into(), Some(TaintLabel::Sink(0)));
    let k2 = graph.add_node(NodeKind::Call, "eval".into(), Some(TaintLabel::Sink(1)));

    graph.add_edge(src, k1, EdgeKind::Argument);
    graph.add_edge(src, k2, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 2);

    let sinks: Vec<_> = findings.iter().map(|f| f.sink).collect();
    assert!(sinks.contains(&0));
    assert!(sinks.contains(&1));
}

#[test]
fn multiple_diamond_pattern() {
    //   src
    //  /   \
    // a     b
    //  \   /
    //   snk
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let a = graph.add_node(NodeKind::Variable, "a".into(), None);
    let b = graph.add_node(NodeKind::Variable, "b".into(), None);
    let snk = graph.add_node(NodeKind::Call, "snk".into(), Some(TaintLabel::Sink(0)));

    graph.add_edge(src, a, EdgeKind::Assignment);
    graph.add_edge(src, b, EdgeKind::Assignment);
    graph.add_edge(a, snk, EdgeKind::Argument);
    graph.add_edge(b, snk, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    // BFS from src: visits src, then a and b. When visiting a, it finds snk and reports one finding.
    // When visiting b, snk is already visited, so no second finding from that path.
    // CPU BFS uses a visited set per source, so only one finding.
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
}

#[test]
fn edge_empty_graph() {
    let graph = TaintGraph::new();
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0);
}

#[test]
fn edge_sources_only() {
    let mut graph = TaintGraph::new();
    let _s1 = graph.add_node(NodeKind::Variable, "s1".into(), Some(TaintLabel::Source(0)));
    let _s2 = graph.add_node(NodeKind::Variable, "s2".into(), Some(TaintLabel::Source(1)));
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0);
}

#[test]
fn edge_sinks_only() {
    let mut graph = TaintGraph::new();
    let _k1 = graph.add_node(NodeKind::Call, "k1".into(), Some(TaintLabel::Sink(0)));
    let _k2 = graph.add_node(NodeKind::Call, "k2".into(), Some(TaintLabel::Sink(1)));
    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0);
}

#[test]
fn edge_cycle_with_source_and_sink() {
    let mut graph = TaintGraph::new();
    let src = graph.add_node(NodeKind::Variable, "src".into(), Some(TaintLabel::Source(0)));
    let a = graph.add_node(NodeKind::Variable, "a".into(), None);
    let b = graph.add_node(NodeKind::Variable, "b".into(), None);
    let snk = graph.add_node(NodeKind::Call, "snk".into(), Some(TaintLabel::Sink(0)));

    // Cycle: src -> a -> b -> a (cycle), and a -> snk
    graph.add_edge(src, a, EdgeKind::Assignment);
    graph.add_edge(a, b, EdgeKind::Assignment);
    graph.add_edge(b, a, EdgeKind::Assignment);
    graph.add_edge(a, snk, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
}

#[test]
fn edge_self_loop_source_sink_node() {
    // A node that is BOTH source and sink with a self-loop
    let mut graph = TaintGraph::new();
    let both = graph.add_node(NodeKind::Variable, "both".into(), Some(TaintLabel::Both(0, 0)));
    graph.add_edge(both, both, EdgeKind::Assignment);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 0, "self-loops without sink arguments must not report");
    let _ = both;
}

#[test]
fn labels_load_credential_theft() {
    let label_set = load_labels(Path::new("rules/malware")).unwrap();

    // Verify process.env is identified as a source (from credential-theft.toml)
    // In merged label set: source[1] = process-env
    assert_eq!(label_node(&label_set, "process.env"), Some(TaintLabel::Source(1)));
    assert_eq!(label_node(&label_set, "foo.process.env.bar"), Some(TaintLabel::Source(1)));

    // Verify fetch is identified as a sink (from credential-theft.toml)
    // In merged label set: sink[4] = http-fetch (pattern: fetch)
    assert_eq!(label_node(&label_set, "fetch"), Some(TaintLabel::Sink(4)));
    assert_eq!(label_node(&label_set, "myfetch"), None);
}

#[test]
fn labels_load_code_execution() {
    let label_set = load_labels(Path::new("rules/malware")).unwrap();

    // Verify eval is identified as a sink (from code-execution.toml)
    // In merged label set: sink[0] = eval (pattern: eval)
    assert_eq!(label_node(&label_set, "eval"), Some(TaintLabel::Sink(0)));
    assert_eq!(label_node(&label_set, "myeval"), None);
}

#[test]
fn malware_process_env_token_to_fetch_evil() {
    let mut graph = TaintGraph::new();
    let process_env = graph.add_node(NodeKind::Variable, "process.env.TOKEN".into(), Some(TaintLabel::Source(0)));
    let token = graph.add_node(NodeKind::Variable, "token".into(), None);
    let url = graph.add_node(NodeKind::Variable, "evil.com".into(), None);
    let fetch_call = graph.add_node(NodeKind::Call, "fetch".into(), Some(TaintLabel::Sink(0)));

    graph.add_edge(process_env, token, EdgeKind::Assignment);
    graph.add_edge(token, fetch_call, EdgeKind::Argument);
    graph.add_edge(url, fetch_call, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
}

#[test]
fn malware_fs_readfile_ssh_to_http_request() {
    let mut graph = TaintGraph::new();
    let fs_read = graph.add_node(NodeKind::Call, "fs.readFile".into(), Some(TaintLabel::Source(0)));
    let ssh_path = graph.add_node(NodeKind::Variable, "~/.ssh/id_rsa".into(), None);
    let data = graph.add_node(NodeKind::Variable, "data".into(), None);
    let http_req = graph.add_node(NodeKind::Call, "http.request".into(), Some(TaintLabel::Sink(0)));

    graph.add_edge(ssh_path, fs_read, EdgeKind::Argument);
    graph.add_edge(fs_read, data, EdgeKind::Return);
    graph.add_edge(data, http_req, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
}

#[test]
fn malware_buffer_from_to_eval() {
    let mut graph = TaintGraph::new();
    let buffer_from = graph.add_node(NodeKind::Call, "Buffer.from".into(), Some(TaintLabel::Source(0)));
    let decoded = graph.add_node(NodeKind::Variable, "decoded".into(), None);
    let eval_call = graph.add_node(NodeKind::Call, "eval".into(), Some(TaintLabel::Sink(0)));

    graph.add_edge(buffer_from, decoded, EdgeKind::Return);
    graph.add_edge(decoded, eval_call, EdgeKind::Argument);

    let findings = analyze(&graph).unwrap();
    assert_eq!(findings.len(), 1);
    assert_eq!(findings[0].source, 0);
    assert_eq!(findings[0].sink, 0);
}