agentshield/analysis/interprocedural/
python.rs

1use std::path::Path;
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5
6use crate::ir::{ScanTarget, SourceLocation};
7
8use super::types::{CALL_EXPR_RE, CallGraph, CallSite, FunctionNode};
9
10pub(crate) static PY_FUNC_DEF_RE: Lazy<Regex> = Lazy::new(|| {
11    Regex::new(r#"(?m)^\s*(?:async\s+)?def\s+([A-Za-z0-9_]+)\s*\(([^)]*)\)"#).expect("valid regex")
12});
13
14pub(crate) fn parse_python_file(
15    graph: &mut CallGraph,
16    file_path: &Path,
17    content: &str,
18    target: &ScanTarget,
19) {
20    let lines: Vec<&str> = content.lines().collect();
21
22    // 1. Extract function definitions
23    for (line_idx, line) in lines.iter().enumerate() {
24        let line_num = line_idx + 1;
25        if let Some(cap) = PY_FUNC_DEF_RE.captures(line) {
26            let func_name = cap[1].to_string();
27            let raw_params = &cap[2];
28
29            let params: Vec<String> = raw_params
30                .split(',')
31                .map(|p| {
32                    p.split(':')
33                        .next()
34                        .unwrap_or("")
35                        .split('=')
36                        .next()
37                        .unwrap_or("")
38                        .trim()
39                        .to_string()
40                })
41                .filter(|p| !p.is_empty() && p != "self" && p != "cls")
42                .collect();
43
44            let start_line = line_num;
45            let func_indent = line.len() - line.trim_start().len();
46            let mut end_line = start_line;
47
48            for (sub_idx, next_line) in lines.iter().enumerate().skip(start_line) {
49                let next_trimmed = next_line.trim();
50                if next_trimmed.is_empty() || next_trimmed.starts_with('#') {
51                    continue;
52                }
53                let next_indent = next_line.len() - next_line.trim_start().len();
54                if next_indent <= func_indent
55                    || next_trimmed.starts_with('@')
56                    || next_trimmed.starts_with("def ")
57                    || next_trimmed.starts_with("async def ")
58                    || next_trimmed.starts_with("class ")
59                {
60                    break;
61                }
62                end_line = sub_idx + 1;
63            }
64
65            let func_sinks = target
66                .data
67                .sinks
68                .iter()
69                .filter(|s| {
70                    s.location.file == file_path
71                        && s.location.line >= start_line
72                        && s.location.line <= end_line
73                })
74                .cloned()
75                .collect();
76
77            let node = FunctionNode {
78                name: func_name.clone(),
79                file_path: file_path.to_path_buf(),
80                params,
81                start_line,
82                end_line,
83                location: SourceLocation {
84                    file: file_path.to_path_buf(),
85                    line: start_line,
86                    column: func_indent,
87                    end_line: Some(end_line),
88                    end_column: None,
89                },
90                sinks: func_sinks,
91            };
92
93            graph.functions.entry(func_name).or_default().push(node);
94        }
95    }
96
97    // 2. Extract call sites
98    for (line_idx, line) in lines.iter().enumerate() {
99        let line_num = line_idx + 1;
100        let trimmed = line.trim();
101        if trimmed.starts_with('#')
102            || trimmed.starts_with("def ")
103            || trimmed.starts_with("async def ")
104        {
105            continue;
106        }
107
108        for cap in CALL_EXPR_RE.captures_iter(line) {
109            let callee_name = cap[1].to_string();
110            let raw_args = &cap[2];
111
112            // Determine caller function containing this line
113            let caller_name = graph.find_enclosing_function(file_path, line_num);
114
115            let args: Vec<String> = raw_args
116                .split(',')
117                .map(|a| a.split('=').next().unwrap_or("").trim().to_string())
118                .filter(|a| !a.is_empty())
119                .collect();
120
121            graph.call_sites.push(CallSite {
122                caller_name,
123                callee_name,
124                file_path: file_path.to_path_buf(),
125                line_number: line_num,
126                args,
127                location: SourceLocation {
128                    file: file_path.to_path_buf(),
129                    line: line_num,
130                    column: line.find(&cap[0]).unwrap_or(0),
131                    end_line: None,
132                    end_column: None,
133                },
134            });
135        }
136    }
137}
agentshield/analysis/interprocedural/python.rs

agentshield/analysis/interprocedural/
python.rs
