agentshield/parser/python/
classify.rs

1use crate::ir::ArgumentSource;
2use crate::ir::SourceLocation;
3use std::path::Path;
4
5/// Classify a call argument string to determine its source.
6pub(crate) fn classify_argument(
7    args_str: &str,
8    param_names: &std::collections::HashSet<String>,
9    sanitized_vars: &std::collections::HashSet<String>,
10) -> ArgumentSource {
11    let first_arg = args_str.split(',').next().unwrap_or("").trim();
12
13    if first_arg.is_empty() {
14        return ArgumentSource::Unknown;
15    }
16
17    // Check if this is a sanitized variable first
18    let ident = first_arg.split('.').next().unwrap_or(first_arg);
19    let ident = ident.split('[').next().unwrap_or(ident);
20    if let Some(sanitizer) = sanitized_label_for_var(ident, sanitized_vars) {
21        return ArgumentSource::Sanitized { sanitizer };
22    }
23
24    // String literal. Single quote tokens can appear when a regex-level parse
25    // sees an incomplete multiline literal; keep those conservative.
26    if let Some(val) = strip_python_string_literal(first_arg) {
27        return ArgumentSource::Literal(val.to_string());
28    }
29
30    // f-string or format
31    if first_arg.starts_with("f\"") || first_arg.starts_with("f'") || first_arg.contains(".format(")
32    {
33        return ArgumentSource::Interpolated;
34    }
35
36    // os.environ / env var
37    if first_arg.contains("os.environ") || first_arg.contains("os.getenv") {
38        return ArgumentSource::EnvVar {
39            name: first_arg.to_string(),
40        };
41    }
42
43    // Known function parameter
44    if param_names.contains(ident) {
45        return ArgumentSource::Parameter {
46            name: ident.to_string(),
47        };
48    }
49
50    ArgumentSource::Unknown
51}
52
53pub(crate) fn strip_python_string_literal(arg: &str) -> Option<&str> {
54    arg.strip_prefix('"')
55        .and_then(|inner| inner.strip_suffix('"'))
56        .or_else(|| {
57            arg.strip_prefix('\'')
58                .and_then(|inner| inner.strip_suffix('\''))
59        })
60}
61
62pub(crate) use crate::parser::common::{sanitized_label_for_var, sanitized_var_marker};
63
64pub(crate) fn loc(file: &Path, line: usize) -> SourceLocation {
65    SourceLocation {
66        file: file.to_path_buf(),
67        line,
68        column: 0,
69        end_line: Some(line),
70        end_column: Some(0),
71    }
72}
73
74pub(crate) fn loc_from_range(
75    file: &Path,
76    line: usize,
77    source_line: &str,
78    start_byte: usize,
79    end_byte: usize,
80) -> SourceLocation {
81    SourceLocation {
82        file: file.to_path_buf(),
83        line,
84        column: source_line[..start_byte].chars().count(),
85        end_line: Some(line),
86        end_column: Some(source_line[..end_byte].chars().count()),
87    }
88}
agentshield/parser/python/classify.rs

agentshield/parser/python/
classify.rs
