agentshield/parser/python/
scanner.rs

1use std::collections::HashSet;
2use std::path::Path;
3
4use crate::analysis::sensitivity::looks_sensitive_name;
5use crate::ir::ArgumentSource;
6use crate::ir::execution_surface::*;
7use crate::parser::{CallSite, ParsedFile};
8
9use super::classify::{classify_argument, loc, loc_from_range};
10use super::patterns::*;
11
12pub(crate) fn scan_python_source(
13    content: &str,
14    file_path: &Path,
15    param_names: &HashSet<String>,
16    http_client_vars: &HashSet<String>,
17    parsed: &mut ParsedFile,
18) {
19    let lines: Vec<&str> = content.lines().collect();
20    let mut current_functions: Vec<(String, usize)> = Vec::new();
21
22    for (line_idx, line) in lines.iter().enumerate() {
23        let line_num = line_idx + 1;
24        let trimmed = line.trim();
25        let indent = line.chars().take_while(|c| c.is_whitespace()).count();
26
27        if !trimmed.is_empty() {
28            while current_functions
29                .last()
30                .is_some_and(|(_, function_indent)| indent <= *function_indent)
31            {
32                current_functions.pop();
33            }
34        }
35        if let Some(cap) = FUNC_DEF_RE.captures(line) {
36            current_functions.push((cap[1].to_string(), indent));
37        }
38
39        // Skip comments
40        if trimmed.starts_with('#') {
41            continue;
42        }
43
44        // A definition header has the same `name(args)` shape as a call
45        // for the regex below. It establishes scope, but is not a call
46        // site and must not participate in cross-file analysis.
47        if FUNC_DEF_RE.is_match(line) {
48            continue;
49        }
50
51        // Check env var access
52        for cap in ENV_ACCESS_RE.captures_iter(line) {
53            let var_name = cap
54                .get(1)
55                .or_else(|| cap.get(2))
56                .or_else(|| cap.get(3))
57                .map(|m| m.as_str().to_string())
58                .unwrap_or_default();
59            let is_sensitive = looks_sensitive_name(&var_name);
60            parsed.env_accesses.push(EnvAccess {
61                var_name: ArgumentSource::Literal(var_name),
62                is_sensitive,
63                location: loc(file_path, line_num),
64            });
65        }
66
67        // Check function calls
68        for cap in CALL_RE.captures_iter(line) {
69            let func_name = &cap[1];
70            let args_str = &cap[2];
71            let call_range = cap.get(0).expect("call capture");
72            let call_location = loc_from_range(
73                file_path,
74                line_num,
75                line,
76                call_range.start(),
77                call_range.end(),
78            );
79
80            let arg_source = classify_argument(args_str, param_names, &parsed.sanitized_vars);
81
82            // Record CallSite for cross-file analysis
83            let all_args = args_str
84                .split(',')
85                .map(|a| classify_argument(a.trim(), param_names, &parsed.sanitized_vars))
86                .collect::<Vec<_>>();
87            parsed.call_sites.push(CallSite {
88                callee: func_name.to_string(),
89                arguments: all_args,
90                caller: current_functions.last().map(|(name, _)| name.clone()),
91                location: call_location.clone(),
92            });
93
94            // Subprocess/command execution
95            if SUBPROCESS_PATTERNS
96                .iter()
97                .any(|p| func_name.ends_with(p) || func_name == *p)
98            {
99                parsed.commands.push(CommandInvocation {
100                    function: func_name.to_string(),
101                    command_arg: arg_source.clone(),
102                    location: call_location.clone(),
103                });
104            }
105
106            // Network operations
107            if NETWORK_PATTERNS
108                .iter()
109                .any(|p| func_name.ends_with(p) || func_name == *p)
110            {
111                let sends_data = func_name.contains("post")
112                    || func_name.contains("put")
113                    || func_name.contains("patch")
114                    || args_str.contains("data=")
115                    || args_str.contains("json=");
116                let method = if func_name.contains("get") {
117                    Some("GET".into())
118                } else if func_name.contains("post") {
119                    Some("POST".into())
120                } else if func_name.contains("put") {
121                    Some("PUT".into())
122                } else {
123                    None
124                };
125                parsed.network_operations.push(NetworkOperation {
126                    function: func_name.to_string(),
127                    url_arg: arg_source.clone(),
128                    method,
129                    sends_data,
130                    location: call_location.clone(),
131                });
132            }
133
134            // Dynamic exec
135            if DYNAMIC_EXEC_PATTERNS.contains(&func_name) {
136                parsed.dynamic_exec.push(DynamicExec {
137                    function: func_name.to_string(),
138                    code_arg: arg_source.clone(),
139                    location: call_location.clone(),
140                });
141            }
142
143            // File operations (open with write mode)
144            if FILE_READ_PATTERNS
145                .iter()
146                .any(|p| func_name.ends_with(p) || func_name == *p)
147            {
148                let op_type = if args_str.contains("'w")
149                    || args_str.contains("\"w")
150                    || args_str.contains("'a")
151                    || args_str.contains("\"a")
152                {
153                    FileOpType::Write
154                } else {
155                    FileOpType::Read
156                };
157                parsed.file_operations.push(FileOperation {
158                    operation: op_type,
159                    path_arg: arg_source.clone(),
160                    location: call_location.clone(),
161                });
162            }
163
164            // HTTP client variable method calls (FN-1 fix):
165            // Detect `client.get(url)` where `client` was bound from
166            // `async with AsyncClient() as client:`.
167            if func_name.contains('.') {
168                let parts: Vec<&str> = func_name.rsplitn(2, '.').collect();
169                if parts.len() == 2 {
170                    let method = parts[0];
171                    let obj = parts[1];
172                    if http_client_vars.contains(obj) && HTTP_CLIENT_METHODS.contains(&method) {
173                        let sends_data = method == "post"
174                            || method == "put"
175                            || method == "patch"
176                            || args_str.contains("data=")
177                            || args_str.contains("json=");
178                        let http_method = match method {
179                            "get" => Some("GET".into()),
180                            "post" => Some("POST".into()),
181                            "put" => Some("PUT".into()),
182                            "delete" => Some("DELETE".into()),
183                            "head" => Some("HEAD".into()),
184                            "patch" => Some("PATCH".into()),
185                            _ => None,
186                        };
187                        parsed.network_operations.push(NetworkOperation {
188                            function: func_name.to_string(),
189                            url_arg: arg_source.clone(),
190                            method: http_method,
191                            sends_data,
192                            location: call_location.clone(),
193                        });
194                    }
195                }
196            }
197        }
198
199        // GitPython command execution (FN-2 fix):
200        // Detect `repo.git.log(...)`, `repo.git.add(...)`, etc.
201        for cap in GITPYTHON_RE.captures_iter(line) {
202            let full_call = format!("{}.git.{}", &cap[1], &cap[2]);
203            let args_str = &cap[3];
204            let arg_source = classify_argument(args_str, param_names, &parsed.sanitized_vars);
205            let call_range = cap.get(0).expect("GitPython call capture");
206            parsed.commands.push(CommandInvocation {
207                function: full_call,
208                command_arg: arg_source,
209                location: loc_from_range(
210                    file_path,
211                    line_num,
212                    line,
213                    call_range.start(),
214                    call_range.end(),
215                ),
216            });
217        }
218
219        // Multi-line call detection: handle calls like
220        //   client.get(
221        //       url,
222        //       follow_redirects=True,
223        //   )
224        // where CALL_RE fails because `(` and `)` are on different lines.
225        if let Some(cap) = PARTIAL_CALL_RE.captures(trimmed) {
226            let func_name = &cap[1];
227            let call_range = cap.get(1).expect("partial call name capture");
228            let trim_offset = line.find(trimmed).unwrap_or_default();
229            let call_location = loc_from_range(
230                file_path,
231                line_num,
232                line,
233                trim_offset + call_range.start(),
234                trim_offset + call_range.end(),
235            );
236
237            let is_http_client_var_call = if func_name.contains('.') {
238                let parts: Vec<&str> = func_name.rsplitn(2, '.').collect();
239                parts.len() == 2
240                    && http_client_vars.contains(parts[1])
241                    && HTTP_CLIENT_METHODS.contains(&parts[0])
242            } else {
243                false
244            };
245
246            let is_network = NETWORK_PATTERNS
247                .iter()
248                .any(|p| func_name.ends_with(p) || func_name == *p)
249                || is_http_client_var_call;
250
251            let is_subprocess = SUBPROCESS_PATTERNS
252                .iter()
253                .any(|p| func_name.ends_with(p) || func_name == *p);
254
255            let is_dynamic = DYNAMIC_EXEC_PATTERNS.contains(&func_name);
256
257            if is_network || is_subprocess || is_dynamic {
258                // Look ahead up to 10 lines to find the first non-kwarg argument
259                let lookahead_limit = (line_idx + 10).min(lines.len());
260                let mut first_arg: Option<ArgumentSource> = None;
261                for future_line in &lines[line_idx + 1..lookahead_limit] {
262                    let trimmed_arg = future_line.trim().trim_end_matches(',');
263                    if trimmed_arg.is_empty() || trimmed_arg.starts_with('#') {
264                        continue;
265                    }
266                    if trimmed_arg == ")" || trimmed_arg.starts_with("):") {
267                        break;
268                    }
269                    if trimmed_arg.contains('=') && !trimmed_arg.starts_with("url=") {
270                        continue;
271                    }
272                    let candidate = if let Some(stripped) = trimmed_arg.strip_prefix("url=") {
273                        stripped.trim()
274                    } else {
275                        trimmed_arg
276                    };
277                    first_arg = Some(classify_argument(
278                        candidate,
279                        param_names,
280                        &parsed.sanitized_vars,
281                    ));
282                    break;
283                }
284
285                let arg_source = first_arg.unwrap_or(ArgumentSource::Unknown);
286
287                if is_network {
288                    let method = if func_name.contains("get") {
289                        Some("GET".into())
290                    } else if func_name.contains("post") {
291                        Some("POST".into())
292                    } else if func_name.contains("put") {
293                        Some("PUT".into())
294                    } else {
295                        None
296                    };
297                    parsed.network_operations.push(NetworkOperation {
298                        function: func_name.to_string(),
299                        url_arg: arg_source,
300                        method,
301                        sends_data: false,
302                        location: call_location,
303                    });
304                } else if is_subprocess {
305                    parsed.commands.push(CommandInvocation {
306                        function: func_name.to_string(),
307                        command_arg: arg_source,
308                        location: call_location,
309                    });
310                } else if is_dynamic {
311                    parsed.dynamic_exec.push(DynamicExec {
312                        function: func_name.to_string(),
313                        code_arg: arg_source,
314                        location: call_location,
315                    });
316                }
317            }
318        }
319    }
320}
agentshield/parser/python/scanner.rs

agentshield/parser/python/
scanner.rs
