agentshield/parser/python/
defs.rs

1use std::collections::HashSet;
2use std::path::Path;
3
4use crate::analysis::cross_file::{SanitizerCategory, sanitizer_category, sanitizer_label};
5use crate::parser::{FunctionDef, FunctionParam, ParsedFile};
6
7use super::classify::{loc, sanitized_var_marker};
8use super::patterns::{FUNC_DEF_RE, HTTP_CLIENT_CTX_RE, SANITIZER_ASSIGN_RE};
9
10pub(crate) fn collect_sanitizer_vars(content: &str, parsed: &mut ParsedFile) {
11    for cap in SANITIZER_ASSIGN_RE.captures_iter(content) {
12        let var_name = &cap[1];
13        let func_name = &cap[2];
14        if sanitizer_category(func_name)
15            .is_some_and(|category| !matches!(category, SanitizerCategory::Redaction))
16        {
17            parsed.sanitized_vars.insert(var_name.to_string());
18            if let Some(label) = sanitizer_label(func_name) {
19                parsed
20                    .sanitized_vars
21                    .insert(sanitized_var_marker(var_name, &label));
22            }
23        }
24    }
25}
26
27pub(crate) fn collect_function_defs_and_params(
28    content: &str,
29    file_path: &Path,
30    parsed: &mut ParsedFile,
31) -> HashSet<String> {
32    let mut param_names = HashSet::new();
33    for cap in FUNC_DEF_RE.captures_iter(content) {
34        let func_name = &cap[1];
35        let params_str = &cap[2];
36        // In Python, functions starting with _ are conventionally private
37        let is_exported = !func_name.starts_with('_');
38        let func_line = content[..cap.get(0).map(|m| m.start()).unwrap_or(0)]
39            .lines()
40            .count()
41            + 1;
42        let function_location = loc(file_path, func_line);
43
44        let mut func_params = Vec::new();
45        for param in params_str.split(',') {
46            let param = param.trim().split(':').next().unwrap_or("").trim();
47            let param = param.split('=').next().unwrap_or("").trim();
48            if !param.is_empty() && param != "self" && param != "cls" {
49                param_names.insert(param.to_string());
50                func_params.push(param.to_string());
51                parsed.function_params.push(FunctionParam {
52                    function_name: func_name.to_string(),
53                    param_name: param.to_string(),
54                    location: function_location.clone(),
55                });
56            }
57        }
58
59        parsed.function_defs.push(FunctionDef {
60            name: func_name.to_string(),
61            params: func_params,
62            is_exported,
63            location: function_location,
64        });
65    }
66    param_names
67}
68
69pub(crate) fn collect_http_client_vars(content: &str) -> HashSet<String> {
70    let mut http_client_vars = HashSet::new();
71    for cap in HTTP_CLIENT_CTX_RE.captures_iter(content) {
72        http_client_vars.insert(cap[1].to_string());
73    }
74    http_client_vars
75}
agentshield/parser/python/defs.rs

agentshield/parser/python/
defs.rs
