agentshield/parser/python/
classify.rs

1use crate::analysis::cross_file::SanitizerCategory;
2use crate::ir::ArgumentSource;
3use crate::ir::SourceLocation;
4use std::path::Path;
5
6/// Classify a call argument string to determine its source.
7pub(crate) fn classify_argument(
8    args_str: &str,
9    param_names: &std::collections::HashSet<String>,
10    sanitized_vars: &std::collections::HashSet<String>,
11) -> ArgumentSource {
12    let first_arg = args_str.split(',').next().unwrap_or("").trim();
13
14    if first_arg.is_empty() {
15        return ArgumentSource::Unknown;
16    }
17
18    // Check if this is a sanitized variable first
19    let ident = first_arg.split('.').next().unwrap_or(first_arg);
20    let ident = ident.split('[').next().unwrap_or(ident);
21    if let Some(sanitizer) = sanitized_label_for_var(ident, sanitized_vars) {
22        return ArgumentSource::Sanitized { sanitizer };
23    }
24
25    // String literal. Single quote tokens can appear when a regex-level parse
26    // sees an incomplete multiline literal; keep those conservative.
27    if let Some(val) = strip_python_string_literal(first_arg) {
28        return ArgumentSource::Literal(val.to_string());
29    }
30
31    // f-string or format
32    if first_arg.starts_with("f\"") || first_arg.starts_with("f'") || first_arg.contains(".format(")
33    {
34        return ArgumentSource::Interpolated;
35    }
36
37    // os.environ / env var
38    if first_arg.contains("os.environ") || first_arg.contains("os.getenv") {
39        return ArgumentSource::EnvVar {
40            name: first_arg.to_string(),
41        };
42    }
43
44    // Known function parameter
45    if param_names.contains(ident) {
46        return ArgumentSource::Parameter {
47            name: ident.to_string(),
48        };
49    }
50
51    ArgumentSource::Unknown
52}
53
54pub(crate) fn strip_python_string_literal(arg: &str) -> Option<&str> {
55    arg.strip_prefix('"')
56        .and_then(|inner| inner.strip_suffix('"'))
57        .or_else(|| {
58            arg.strip_prefix('\'')
59                .and_then(|inner| inner.strip_suffix('\''))
60        })
61}
62
63pub(crate) fn sanitized_var_marker(var_name: &str, sanitizer_label: &str) -> String {
64    format!("{var_name}::{sanitizer_label}")
65}
66
67pub(crate) fn sanitized_label_for_var(
68    ident: &str,
69    sanitized_vars: &std::collections::HashSet<String>,
70) -> Option<String> {
71    for category in [
72        SanitizerCategory::Path,
73        SanitizerCategory::Network,
74        SanitizerCategory::TypeCoercion,
75    ] {
76        let prefix = format!("{ident}::{}:", category.as_str());
77        let mut matches: Vec<&str> = sanitized_vars
78            .iter()
79            .filter(|value| value.starts_with(&prefix))
80            .map(|s| s.as_str())
81            .collect();
82        matches.sort();
83        if let Some(marker) = matches.first() {
84            return marker.split_once("::").map(|(_, label)| label.to_string());
85        }
86    }
87
88    sanitized_vars.contains(ident).then(|| ident.to_string())
89}
90
91pub(crate) fn loc(file: &Path, line: usize) -> SourceLocation {
92    SourceLocation {
93        file: file.to_path_buf(),
94        line,
95        column: 0,
96        end_line: Some(line),
97        end_column: Some(0),
98    }
99}
100
101pub(crate) fn loc_from_range(
102    file: &Path,
103    line: usize,
104    source_line: &str,
105    start_byte: usize,
106    end_byte: usize,
107) -> SourceLocation {
108    SourceLocation {
109        file: file.to_path_buf(),
110        line,
111        column: source_line[..start_byte].chars().count(),
112        end_line: Some(line),
113        end_column: Some(source_line[..end_byte].chars().count()),
114    }
115}
agentshield/parser/python/classify.rs

agentshield/parser/python/
classify.rs
