From f8161ebf402c6fc0af00bd55d77e6725ee1fd90d Mon Sep 17 00:00:00 2001
From: Mathieu Kniewallner <mathieu.kniewallner@gmail.com>
Date: Mon, 25 Sep 2023 20:27:53 +0200
Subject: [PATCH] fix(plugins/B507): also detect class instances

`paramiko` supports passing both a class and a class instance for the
policy in `set_missing_host_key_policy`
(https://github.com/paramiko/paramiko/blob/8e389c77660c5cdae3069b478665427d23012853/paramiko/client.py#L171-L191).
This updates B507 to account for both styles.
---
 bandit/plugins/ssh_no_host_key_verification.py | 16 ++++++++++++----
 examples/no_host_key_verification.py           |  2 ++
 tests/functional/test_functional.py            |  4 ++--
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/bandit/plugins/ssh_no_host_key_verification.py b/bandit/plugins/ssh_no_host_key_verification.py
index 2f4390320..e8edaf93a 100644
--- a/bandit/plugins/ssh_no_host_key_verification.py
+++ b/bandit/plugins/ssh_no_host_key_verification.py
@@ -35,6 +35,8 @@
     CWE information added
 
 """
+import ast
+
 import bandit
 from bandit.core import issue
 from bandit.core import test_properties as test
@@ -46,11 +48,17 @@ def ssh_no_host_key_verification(context):
     if (
         context.is_module_imported_like("paramiko")
         and context.call_function_name == "set_missing_host_key_policy"
+        and context.node.args
     ):
-        if context.call_args and context.call_args[0] in [
-            "AutoAddPolicy",
-            "WarningPolicy",
-        ]:
+        policy_argument = context.node.args[0]
+
+        policy_argument_value = None
+        if isinstance(policy_argument, ast.Attribute):
+            policy_argument_value = policy_argument.attr
+        elif isinstance(policy_argument, ast.Call):
+            policy_argument_value = policy_argument.func.attr
+
+        if policy_argument_value in ["AutoAddPolicy", "WarningPolicy"]:
             return bandit.Issue(
                 severity=bandit.HIGH,
                 confidence=bandit.MEDIUM,
diff --git a/examples/no_host_key_verification.py b/examples/no_host_key_verification.py
index ff12b343e..2e092fe4d 100644
--- a/examples/no_host_key_verification.py
+++ b/examples/no_host_key_verification.py
@@ -3,3 +3,5 @@
 ssh_client = client.SSHClient()
 ssh_client.set_missing_host_key_policy(client.AutoAddPolicy)
 ssh_client.set_missing_host_key_policy(client.WarningPolicy)
+ssh_client.set_missing_host_key_policy(client.AutoAddPolicy())
+ssh_client.set_missing_host_key_policy(client.WarningPolicy())
diff --git a/tests/functional/test_functional.py b/tests/functional/test_functional.py
index 7835e7488..6917462dd 100644
--- a/tests/functional/test_functional.py
+++ b/tests/functional/test_functional.py
@@ -543,8 +543,8 @@ def test_yaml(self):
     def test_host_key_verification(self):
         """Test for ignoring host key verification."""
         expect = {
-            "SEVERITY": {"UNDEFINED": 0, "LOW": 0, "MEDIUM": 0, "HIGH": 2},
-            "CONFIDENCE": {"UNDEFINED": 0, "LOW": 0, "MEDIUM": 2, "HIGH": 0},
+            "SEVERITY": {"UNDEFINED": 0, "LOW": 0, "MEDIUM": 0, "HIGH": 4},
+            "CONFIDENCE": {"UNDEFINED": 0, "LOW": 0, "MEDIUM": 4, "HIGH": 0},
         }
         self.check_example("no_host_key_verification.py", expect)