From e28faaa14f32a20e8f6e702ef5a65107bbcce4c7 Mon Sep 17 00:00:00 2001
From: ROSPL07 <rosplvesit07@gmail.com>
Date: Tue, 16 Sep 2025 16:48:27 +0530
Subject: [PATCH 1/2] feat(plugin): add B901 check for unsafe
 yaml.load/full_load usage

---
 bandit/plugins/yaml_unsafe_loader.py | 58 ++++++++++++++++++++++++++++
 setup.cfg                            |  3 ++
 2 files changed, 61 insertions(+)
 create mode 100644 bandit/plugins/yaml_unsafe_loader.py

diff --git a/bandit/plugins/yaml_unsafe_loader.py b/bandit/plugins/yaml_unsafe_loader.py
new file mode 100644
index 000000000..fc9b193dd
--- /dev/null
+++ b/bandit/plugins/yaml_unsafe_loader.py
@@ -0,0 +1,58 @@
+# bandit/plugins/yaml_unsafe_loader.py
+"""
+B901: Detect usage of yaml.full_load, yaml.unsafe_load, or yaml.load with unsafe Loader.
+Small Bandit plugin to explicitly flag full_load/unsafe_load and non-safe Loader usage.
+"""
+
+import bandit
+from bandit.core import test_properties as test
+
+# Register as a Call-check with a test id so Bandit will run it
+@test.test_id("B901")
+@test.checks("Call")
+def yaml_unsafe_loader(context):
+    """
+    Flag:
+      - yaml.full_load(...)
+      - yaml.unsafe_load(...)
+      - yaml.load(..., Loader=SomeLoader) if loader name does not indicate 'safe'
+    """
+    func_name = context.call_function_name_qual or ""
+
+    # direct unsafe calls: yaml.full_load(), yaml.unsafe_load()
+    if "yaml.full_load" in func_name or "yaml.unsafe_load" in func_name:
+        return bandit.Issue(
+            severity=bandit.MEDIUM,
+            confidence=bandit.HIGH,
+            text="Use of yaml.full_load()/yaml.unsafe_load() detected. Prefer yaml.safe_load() "
+                 "when parsing untrusted YAML."
+        )
+
+    # yaml.load(...) with Loader=... where loader isn't safe
+    if "yaml.load" in func_name:
+        node = context.node  # AST Call node
+        for kw in getattr(node, "keywords", []):
+            if kw.arg == "Loader":
+                val = kw.value
+                loader_name = ""
+                # e.g. Loader=yaml.FullLoader  -> val.attr == 'FullLoader'
+                if hasattr(val, "attr"):
+                    loader_name = val.attr
+                # e.g. Loader=FullLoader  -> val.id == 'FullLoader'
+                elif hasattr(val, "id"):
+                    loader_name = val.id
+                # Fallback: try string literal
+                elif hasattr(val, "s"):
+                    loader_name = getattr(val, "s", "")
+                if loader_name:
+                    # if loader name doesn't contain "safe", flag it
+                    if "safe" not in loader_name.lower():
+                        return bandit.Issue(
+                            severity=bandit.MEDIUM,
+                            confidence=bandit.MEDIUM,
+                            text=f"yaml.load() used with Loader={loader_name!r} - this can be unsafe "
+                                 "for untrusted input. Consider yaml.safe_load()."
+                        )
+
+    # no issue found
+    return None
diff --git a/setup.cfg b/setup.cfg
index e1bab07f5..6aff9f8af 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -166,6 +166,9 @@ bandit.plugins =
     # bandit/plugins/huggingface_unsafe_download.py
     huggingface_unsafe_download = bandit.plugins.huggingface_unsafe_download:huggingface_unsafe_download
 
+    # bandit/plugins/yaml_unsafe_loader.py
+    yaml_unsafe_loader = bandit.plugins.yaml_unsafe_loader:yaml_unsafe_loader
+
 [build_sphinx]
 all_files = 1
 build-dir = doc/build

From 0e997abae4087a08b6f21b0706628e010aa30e11 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 16 Sep 2025 11:20:17 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 bandit/plugins/yaml_unsafe_loader.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/bandit/plugins/yaml_unsafe_loader.py b/bandit/plugins/yaml_unsafe_loader.py
index fc9b193dd..ef8a5220b 100644
--- a/bandit/plugins/yaml_unsafe_loader.py
+++ b/bandit/plugins/yaml_unsafe_loader.py
@@ -3,10 +3,10 @@
 B901: Detect usage of yaml.full_load, yaml.unsafe_load, or yaml.load with unsafe Loader.
 Small Bandit plugin to explicitly flag full_load/unsafe_load and non-safe Loader usage.
 """
-
 import bandit
 from bandit.core import test_properties as test
 
+
 # Register as a Call-check with a test id so Bandit will run it
 @test.test_id("B901")
 @test.checks("Call")
@@ -25,7 +25,7 @@ def yaml_unsafe_loader(context):
             severity=bandit.MEDIUM,
             confidence=bandit.HIGH,
             text="Use of yaml.full_load()/yaml.unsafe_load() detected. Prefer yaml.safe_load() "
-                 "when parsing untrusted YAML."
+            "when parsing untrusted YAML.",
         )
 
     # yaml.load(...) with Loader=... where loader isn't safe
@@ -51,7 +51,7 @@ def yaml_unsafe_loader(context):
                             severity=bandit.MEDIUM,
                             confidence=bandit.MEDIUM,
                             text=f"yaml.load() used with Loader={loader_name!r} - this can be unsafe "
-                                 "for untrusted input. Consider yaml.safe_load()."
+                            "for untrusted input. Consider yaml.safe_load().",
                         )
 
     # no issue found