[`flake8-bandit`] Add Rule for `S506` (unsafe use of yaml load) (#1664)

See: https://github.com/charliermarsh/ruff/issues/1646.

Co-authored-by: Charlie Marsh <charlie.r.marsh@gmail.com>

README.md

@@ -767,11 +767,12 @@ For more, see [flake8-bandit](https://pypi.org/project/flake8-bandit/4.1.1/) on
 | S102 | ExecUsed | Use of `exec` detected |  |
 | S103 | BadFilePermissions | `os.chmod` setting a permissive mask `0o777` on file or directory |  |
 | S104 | HardcodedBindAllInterfaces | Possible binding to all interfaces |  |
-| S105 | HardcodedPasswordString | Possible hardcoded password: `"..."` |  |
+| S105 | HardcodedPasswordString | Possible hardcoded password: "..." |  |
-| S106 | HardcodedPasswordFuncArg | Possible hardcoded password: `"..."` |  |
+| S106 | HardcodedPasswordFuncArg | Possible hardcoded password: "..." |  |
-| S107 | HardcodedPasswordDefault | Possible hardcoded password: `"..."` |  |
+| S107 | HardcodedPasswordDefault | Possible hardcoded password: "..." |  |
-| S108 | HardcodedTempFile | Probable insecure usage of temp file/directory: `"..."` |  |
+| S108 | HardcodedTempFile | Probable insecure usage of temporary file or directory: "..." |  |
-| S324 | HashlibInsecureHashFunction | Probable use of insecure hash functions in hashlib: `"..."` |  |
+| S324 | HashlibInsecureHashFunction | Probable use of insecure hash functions in `hashlib`: "..." |  |
+| S506 | UnsafeYAMLLoad | Probable use of unsafe `yaml.load`. Allows instantiation of arbitrary objects. Consider `yaml.safe_load`. |  |
 
 ### flake8-blind-except (BLE)
 

resources/test/fixtures/flake8_bandit/S506.py

@@ -0,0 +1,31 @@
+import json
+import yaml
+from yaml import CSafeLoader
+from yaml import SafeLoader
+from yaml import SafeLoader as NewSafeLoader
+
+
+def test_yaml_load():
+    ystr = yaml.dump({"a": 1, "b": 2, "c": 3})
+    y = yaml.load(ystr)
+    yaml.dump(y)
+    try:
+        y = yaml.load(ystr, Loader=yaml.CSafeLoader)
+    except AttributeError:
+        # CSafeLoader only exists if you build yaml with LibYAML
+        y = yaml.load(ystr, Loader=yaml.SafeLoader)
+
+
+def test_json_load():
+    # no issue should be found
+    j = json.load("{}")
+
+
+yaml.load("{}", Loader=yaml.Loader)
+
+# no issue should be found
+yaml.load("{}", SafeLoader)
+yaml.load("{}", yaml.SafeLoader)
+yaml.load("{}", CSafeLoader)
+yaml.load("{}", yaml.CSafeLoader)
+yaml.load("{}", NewSafeLoader)

ruff.schema.json

@@ -930,6 +930,9 @@
         "S3",
         "S32",
         "S324",
+        "S5",
+        "S50",
+        "S506",
         "SIM",
         "SIM1",
         "SIM10",

src/checkers/ast.rs

@@ -1896,6 +1896,17 @@ where
                         self.add_check(check);
                     }
                 }
+                if self.settings.enabled.contains(&CheckCode::S506) {
+                    if let Some(check) = flake8_bandit::checks::unsafe_yaml_load(
+                        func,
+                        args,
+                        keywords,
+                        &self.from_imports,
+                        &self.import_aliases,
+                    ) {
+                        self.add_check(check);
+                    }
+                }
                 if self.settings.enabled.contains(&CheckCode::S106) {
                     self.add_checks(
                         flake8_bandit::checks::hardcoded_password_func_arg(keywords).into_iter(),

src/flake8_bandit/checks/mod.rs

@@ -9,6 +9,7 @@ pub use hardcoded_password_string::{
 };
 pub use hardcoded_tmp_directory::hardcoded_tmp_directory;
 pub use hashlib_insecure_hash_functions::hashlib_insecure_hash_functions;
+pub use unsafe_yaml_load::unsafe_yaml_load;
 
 mod assert_used;
 mod bad_file_permissions;
@@ -19,3 +20,4 @@ mod hardcoded_password_func_arg;
 mod hardcoded_password_string;
 mod hardcoded_tmp_directory;
 mod hashlib_insecure_hash_functions;
+mod unsafe_yaml_load;

src/flake8_bandit/checks/unsafe_yaml_load.rs

@@ -0,0 +1,50 @@
+use rustc_hash::{FxHashMap, FxHashSet};
+use rustpython_ast::{Expr, ExprKind, Keyword};
+
+use crate::ast::helpers::{match_module_member, SimpleCallArgs};
+use crate::ast::types::Range;
+use crate::registry::{Check, CheckKind};
+
+/// S506
+pub fn unsafe_yaml_load(
+    func: &Expr,
+    args: &[Expr],
+    keywords: &[Keyword],
+    from_imports: &FxHashMap<&str, FxHashSet<&str>>,
+    import_aliases: &FxHashMap<&str, &str>,
+) -> Option<Check> {
+    if match_module_member(func, "yaml", "load", from_imports, import_aliases) {
+        let call_args = SimpleCallArgs::new(args, keywords);
+        if let Some(loader_arg) = call_args.get_argument("Loader", Some(1)) {
+            if !match_module_member(
+                loader_arg,
+                "yaml",
+                "SafeLoader",
+                from_imports,
+                import_aliases,
+            ) && !match_module_member(
+                loader_arg,
+                "yaml",
+                "CSafeLoader",
+                from_imports,
+                import_aliases,
+            ) {
+                let loader = match &loader_arg.node {
+                    ExprKind::Attribute { attr, .. } => Some(attr.to_string()),
+                    ExprKind::Name { id, .. } => Some(id.to_string()),
+                    _ => None,
+                };
+                return Some(Check::new(
+                    CheckKind::UnsafeYAMLLoad(loader),
+                    Range::from_located(loader_arg),
+                ));
+            }
+        } else {
+            return Some(Check::new(
+                CheckKind::UnsafeYAMLLoad(None),
+                Range::from_located(func),
+            ));
+        }
+    }
+    None
+}

src/flake8_bandit/mod.rs

@@ -22,6 +22,7 @@ mod tests {
     #[test_case(CheckCode::S107, Path::new("S107.py"); "S107")]
     #[test_case(CheckCode::S108, Path::new("S108.py"); "S108")]
     #[test_case(CheckCode::S324, Path::new("S324.py"); "S324")]
+    #[test_case(CheckCode::S506, Path::new("S506.py"); "S506")]
     fn checks(check_code: CheckCode, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", check_code.as_ref(), path.to_string_lossy());
         let checks = test_path(

src/flake8_bandit/snapshots/ruff__flake8_bandit__tests__S506_S506.py.snap

@@ -0,0 +1,25 @@
+---
+source: src/flake8_bandit/mod.rs
+expression: checks
+---
+- kind:
+    UnsafeYAMLLoad: ~
+  location:
+    row: 10
+    column: 8
+  end_location:
+    row: 10
+    column: 17
+  fix: ~
+  parent: ~
+- kind:
+    UnsafeYAMLLoad: Loader
+  location:
+    row: 24
+    column: 23
+  end_location:
+    row: 24
+    column: 34
+  fix: ~
+  parent: ~
+

src/registry.rs

@@ -335,6 +335,7 @@ pub enum CheckCode {
     S107,
     S108,
     S324,
+    S506,
     // flake8-boolean-trap
     FBT001,
     FBT002,
@@ -1075,6 +1076,7 @@ pub enum CheckKind {
     HardcodedPasswordDefault(String),
     HardcodedTempFile(String),
     HashlibInsecureHashFunction(String),
+    UnsafeYAMLLoad(Option<String>),
     // mccabe
     FunctionIsTooComplex(String, usize),
     // flake8-boolean-trap
@@ -1538,6 +1540,7 @@ impl CheckCode {
             CheckCode::S107 => CheckKind::HardcodedPasswordDefault("...".to_string()),
             CheckCode::S108 => CheckKind::HardcodedTempFile("...".to_string()),
             CheckCode::S324 => CheckKind::HashlibInsecureHashFunction("...".to_string()),
+            CheckCode::S506 => CheckKind::UnsafeYAMLLoad(None),
             // mccabe
             CheckCode::C901 => CheckKind::FunctionIsTooComplex("...".to_string(), 10),
             // flake8-boolean-trap
@@ -1941,6 +1944,7 @@ impl CheckCode {
             CheckCode::S107 => CheckCategory::Flake8Bandit,
             CheckCode::S108 => CheckCategory::Flake8Bandit,
             CheckCode::S324 => CheckCategory::Flake8Bandit,
+            CheckCode::S506 => CheckCategory::Flake8Bandit,
             // flake8-simplify
             CheckCode::SIM102 => CheckCategory::Flake8Simplify,
             CheckCode::SIM105 => CheckCategory::Flake8Simplify,
@@ -2317,6 +2321,7 @@ impl CheckKind {
             CheckKind::HardcodedPasswordDefault(..) => &CheckCode::S107,
             CheckKind::HardcodedTempFile(..) => &CheckCode::S108,
             CheckKind::HashlibInsecureHashFunction(..) => &CheckCode::S324,
+            CheckKind::UnsafeYAMLLoad(..) => &CheckCode::S506,
             // mccabe
             CheckKind::FunctionIsTooComplex(..) => &CheckCode::C901,
             // flake8-boolean-trap
@@ -3260,23 +3265,31 @@ impl CheckKind {
             CheckKind::HardcodedPasswordString(string)
             | CheckKind::HardcodedPasswordFuncArg(string)
             | CheckKind::HardcodedPasswordDefault(string) => {
-                format!(
+                format!("Possible hardcoded password: \"{}\"", string.escape_debug())
-                    "Possible hardcoded password: `\"{}\"`",
-                    string.escape_debug()
-                )
             }
             CheckKind::HardcodedTempFile(string) => {
                 format!(
-                    "Probable insecure usage of temp file/directory: `\"{}\"`",
+                    "Probable insecure usage of temporary file or directory: \"{}\"",
                     string.escape_debug()
                 )
             }
             CheckKind::HashlibInsecureHashFunction(string) => {
                 format!(
-                    "Probable use of insecure hash functions in hashlib: `\"{}\"`",
+                    "Probable use of insecure hash functions in `hashlib`: \"{}\"",
                     string.escape_debug()
                 )
             }
+            CheckKind::UnsafeYAMLLoad(loader) => match loader {
+                Some(name) => {
+                    format!(
+                        "Probable use of unsafe loader `{name}` with `yaml.load`. Allows \
+                         instantiation of arbitrary objects. Consider `yaml.safe_load`."
+                    )
+                }
+                None => "Probable use of unsafe `yaml.load`. Allows instantiation of arbitrary \
+                         objects. Consider `yaml.safe_load`."
+                    .to_string(),
+            },
             // flake8-blind-except
             CheckKind::BlindExcept(name) => format!("Do not catch blind exception: `{name}`"),
             // mccabe