Unit test for --custom-waf-bypass-payload option

Author: axel3rd

log4j-scan.py

@@ -315,11 +315,11 @@ def scan_url(url, callback_host, proxies, args):
     payload = '${jndi:ldap://%s/%s}' % (host_def, random_string)
     payloads = [payload]
     if args.waf_bypass_payloads:
-        payloads.extend(generate_waf_bypass_payloads(f'{parsed_url["host"]}.{callback_host}', random_string))
+        payloads.extend(generate_waf_bypass_payloads(host_def, random_string))
 
     if args.cve_2021_45046:
         cprint(f"[•] Scanning for CVE-2021-45046 (Log4j v2.15.0 Patch Bypass - RCE)", "yellow")
-        payloads = get_cve_2021_45046_payloads(f'{parsed_url["host"]}.{callback_host}', random_string)
+        payloads = get_cve_2021_45046_payloads(host_def, random_string)
 
     auth = None
     if args.basic_auth_user:
@@ -386,7 +386,7 @@ def main(options):
     if args.proxy:
         proxies = {"http": args.proxy, "https": args.proxy}
 
-    if args.custom_waf_bypass_payload:
+    if args.custom_waf_bypass_payload and args.custom_waf_bypass_payload not in waf_bypass_payloads:
         waf_bypass_payloads.append(args.custom_waf_bypass_payload)
 
     urls = []

tests/test_log4j_scan.py

@@ -28,12 +28,30 @@ def test_default(requests_mock, capsys):
     assert adapter_endpoint.call_count == 1
     assert adapter_dns_save.call_count == 1
     assert '.interact.sh/' in captured.out
-    assert 'Targets does not seem to be vulnerable' in captured.out
+    assert 'Targets do not seem to be vulnerable' in captured.out
     assert 'jndi' in adapter_endpoint.last_request.url
     assert re.match(r'\${jndi:ldap://localhost\..*.interact\.sh/.*}', adapter_endpoint.last_request.headers['User-Agent'])
     assert 'Authorization' not in adapter_endpoint.last_request.headers
 
 
+def test_custom_waf_bypass_payload_custom_dns_callback_host(requests_mock):
+    adapter_endpoint = requests_mock.get(LOCALHOST)
+    
+    log4j_scan.main(['-u', LOCALHOST, '--custom-dns-callback-host', DNS_CUSTOM , '--waf-bypass', '--custom-waf-bypass-payload', 'test://{{callback_host}}/{{random}}' ])    
+
+    assert adapter_endpoint.call_count == 25
+    assert re.match(r'test://localhost\.custom.dns.callback/.*', adapter_endpoint.request_history[24].headers['User-Agent'])
+
+
+def test_custom_waf_bypass_payload_custom_tcp_callback_host(requests_mock):
+    adapter_endpoint = requests_mock.get(LOCALHOST)
+
+    log4j_scan.main(['-u', LOCALHOST, '--custom-tcp-callback-host', '10.42.42.42:80' , '--waf-bypass', '--custom-waf-bypass-payload', 'test://{{callback_host}}/{{random}}' ])    
+
+    assert adapter_endpoint.call_count == 25
+    assert re.match(r'test://10\.42\.42\.42:80/.*', adapter_endpoint.request_history[24].headers['User-Agent'])
+
+
 def test_custom_dns_callback_host(requests_mock, capsys):
     adapter_endpoint = requests_mock.get(LOCALHOST)