pyrograph 0.1.0

GPU-accelerated taint analysis for supply chain malware detection
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

//! Go parser entry points.

use std::collections::HashMap;

use crate::error::Result;
use crate::ir::{EdgeKind, NodeId, NodeKind, TaintGraph};
use crate::labels::LabelSet;
use tree_sitter::Parser;

mod labels_default;
mod visitor;

pub use labels_default::default_label_set;
use labels_default::apply_labels;

/// Parser context for Go AST traversal.
pub(crate) struct GoParser {
    pub source: String,
    pub graph: TaintGraph,
    pub scopes: Vec<HashMap<String, NodeId>>,
    pub imports: HashMap<String, String>,
    pub current_function: Option<NodeId>,
    pub current_param_index: HashMap<NodeId, Vec<NodeId>>,
    pub command_vars: HashMap<NodeId, bool>,
    pub reflect_vars: HashMap<NodeId, bool>,
}

impl GoParser {
    pub(crate) fn new() -> Self {
        Self {
            source: String::new(),
            graph: TaintGraph::new(),
            scopes: vec![HashMap::new()],
            imports: HashMap::new(),
            current_function: None,
            current_param_index: HashMap::new(),
            command_vars: HashMap::new(),
            reflect_vars: HashMap::new(),
        }
    }

    pub(crate) fn parse_source(&mut self, source: &str) -> Result<()> {
        let mut parser = Parser::new();
        parser
            .set_language(&tree_sitter_go::LANGUAGE.into())
            .map_err(|error| crate::error::Error::Analysis(format!("failed to configure Go parser language: {error}")))?;

        self.source = source.to_string();
        let tree = parser
            .parse(source, None)
            .ok_or_else(|| crate::error::Error::Analysis("failed to parse Go source".to_string()))?;

        self.walk_node(tree.root_node());
        Ok(())
    }
}

/// Parse a single Go source file into a taint graph.
///
/// Custom labels can be supplied for unit tests and alternate rule sets.
pub fn parse_go_with_labels(
    source: &str,
    _filename: &str,
    labels: Option<&LabelSet>,
) -> Result<TaintGraph> {
    let mut parser = GoParser::new();
    parser.parse_source(source)?;

    let mut graph = parser.graph;
    let label_set = labels.cloned().unwrap_or_else(default_label_set);
    apply_labels(&mut graph, &label_set);
    graph.set_label_set(label_set);
    Ok(graph)
}

/// Parse a single Go source file into a taint graph.
pub fn parse_go(source: &str, filename: &str) -> Result<TaintGraph> {
    parse_go_with_labels(source, filename, None)
}