agentshield/fix/
deserializer.rs

1use once_cell::sync::Lazy;
2use regex::Regex;
3use std::path::Path;
4
5use super::{AppliedFix, FixOutput};
6
7static LOADER_ARG_RE: Lazy<Regex> =
8    Lazy::new(|| Regex::new(r"Loader=[A-Za-z0-9_.]+").expect("valid regex"));
9
10/// Fix unsafe deserializer patterns in Python source code (SHIELD-016).
11pub fn fix_unsafe_deserializers(content: &str, path: &Path) -> Option<FixOutput> {
12    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
13    if ext != "py" {
14        return None;
15    }
16
17    let mut modified_lines = Vec::new();
18    let mut fixes = Vec::new();
19    let mut made_changes = false;
20
21    let lines: Vec<&str> = content.lines().collect();
22
23    for (line_idx, line) in lines.iter().enumerate() {
24        let trimmed = line.trim();
25
26        // 1. Rewrite yaml.load to yaml.safe_load or add SafeLoader
27        if trimmed.contains("yaml.load(")
28            && !trimmed.contains("yaml.safe_load(")
29            && !trimmed.contains("SafeLoader")
30            && !trimmed.contains("CSafeLoader")
31            && !trimmed.starts_with('#')
32        {
33            let new_line = if line.contains("Loader=") {
34                LOADER_ARG_RE
35                    .replace(line, "Loader=yaml.SafeLoader")
36                    .to_string()
37            } else if line.contains("yaml.load(") {
38                // If it's a simple yaml.load(data), replace with yaml.safe_load(data)
39                line.replace("yaml.load(", "yaml.safe_load(")
40            } else {
41                line.to_string()
42            };
43
44            if new_line != *line {
45                fixes.push(AppliedFix {
46                    rule_id: "SHIELD-016".into(),
47                    description: "Replaced unsafe 'yaml.load' with 'yaml.safe_load'".into(),
48                    line_number: line_idx + 1,
49                });
50                modified_lines.push(new_line);
51                made_changes = true;
52                continue;
53            }
54        }
55
56        // 2. Rewrite pickle.loads to json.loads
57        if trimmed.contains("pickle.loads(") && !trimmed.starts_with('#') {
58            let new_line = line.replace("pickle.loads(", "json.loads(");
59            fixes.push(AppliedFix {
60                rule_id: "SHIELD-016".into(),
61                description: "Replaced insecure 'pickle.loads' with 'json.loads'".into(),
62                line_number: line_idx + 1,
63            });
64            modified_lines.push(new_line);
65            made_changes = true;
66            continue;
67        }
68
69        modified_lines.push(line.to_string());
70    }
71
72    // Check if we need to update `import pickle` to `import json`
73    if made_changes
74        && content.contains("import pickle")
75        && !modified_lines.join("\n").contains("pickle.")
76    {
77        for line in &mut modified_lines {
78            if line.trim() == "import pickle" {
79                *line = line.replace("import pickle", "import json");
80            }
81        }
82    }
83
84    if !made_changes {
85        return None;
86    }
87
88    let mut output = modified_lines.join("\n");
89    if content.ends_with('\n') {
90        output.push('\n');
91    }
92
93    Some(FixOutput {
94        content: output,
95        fixes,
96    })
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn test_fix_yaml_load_to_safe_load() {
105        let code = "import yaml\ndata = yaml.load(user_input)\n";
106        let res = fix_unsafe_deserializers(code, Path::new("server.py")).unwrap();
107        assert_eq!(
108            res.content,
109            "import yaml\ndata = yaml.safe_load(user_input)\n"
110        );
111        assert_eq!(res.fixes.len(), 1);
112        assert_eq!(res.fixes[0].rule_id, "SHIELD-016");
113    }
114
115    #[test]
116    fn test_fix_pickle_loads_to_json_loads() {
117        let code = "import pickle\ndata = pickle.loads(payload)\n";
118        let res = fix_unsafe_deserializers(code, Path::new("loader.py")).unwrap();
119        assert_eq!(res.content, "import json\ndata = json.loads(payload)\n");
120        assert_eq!(res.fixes.len(), 1);
121    }
122}
agentshield/fix/deserializer.rs

agentshield/fix/
deserializer.rs
