Call the get_hostname_options in find_hostname_in_context instead

Author: bitterpanda63

aikido_zen/vulnerabilities/ssrf/find_hostname_in_context.py

@@ -4,6 +4,7 @@
 
 from aikido_zen.helpers.extract_strings_from_context import extract_strings_from_context
 from .find_hostname_in_userinput import find_hostname_in_userinput
+from .get_hostname_options import get_hostname_options
 from .is_request_to_itself import is_request_to_itself
 from ...context import Context
 
@@ -19,8 +20,11 @@ def find_hostname_in_context(hostname, context: Context, port):
     if is_request_to_itself(context.url, hostname, port):
         return None
 
+    # Gets the different hostname options: with/without punycode, with/without brackets for IPv6
+    hostname_options = get_hostname_options(hostname)
+
     for user_input, path, source in extract_strings_from_context(context):
-        found = find_hostname_in_userinput(user_input, hostname, port)
+        found = find_hostname_in_userinput(user_input, hostname_options, port)
         if found:
             return {
                 "source": source,

aikido_zen/vulnerabilities/ssrf/find_hostname_in_userinput.py

@@ -3,19 +3,16 @@
 """
 
 from typing import List
-
 from aikido_zen.helpers.get_port_from_url import get_port_from_url
 from aikido_zen.helpers.try_parse_url import try_parse_url
 
 
-def find_hostname_in_userinput(user_input, hostname, port=None):
+def find_hostname_in_userinput(user_input, hostname_options: List[str], port=None):
     """
     Returns true if the hostname is in userinput
     """
     if len(user_input) <= 1:
         return False
-
-    hostname_options = get_hostname_options(hostname)
     if len(hostname_options) == 0:
         return False
 

aikido_zen/vulnerabilities/ssrf/find_hostname_in_userinput_test.py

@@ -1,5 +1,13 @@
 import pytest
-from .find_hostname_in_userinput import find_hostname_in_userinput
+from .find_hostname_in_userinput import (
+    find_hostname_in_userinput as _find_hostname_in_userinput,
+)
+from .get_hostname_options import get_hostname_options
+
+
+def find_hostname_in_userinput(user_input, hostname, port=None):
+    hostname_options = get_hostname_options(hostname)
+    return _find_hostname_in_userinput(user_input, hostname_options, port)
 
 
 def test_returns_false_if_user_input_and_hostname_are_empty():