run same urls for urllib3 as for requests

Author: bitterpanda63

aikido_zen/sinks/tests/requests_test.py …sinks/tests/requests_and_urrlib3_test.pyaikido_zen/sinks/tests/requests_test.py renamed to aikido_zen/sinks/tests/requests_and_urrlib3_test.py aikido_zen/sinks/tests/requests_test.py renamed to aikido_zen/sinks/tests/requests_and_urrlib3_test.py

@@ -1,12 +1,14 @@
 import os
 import pytest
+
 from aikido_zen.context import Context, current_context
 from aikido_zen.thread.thread_cache import ThreadCache, get_cache
 from aikido_zen.errors import AikidoSSRF
 from aikido_zen.background_process.comms import reset_comms
 import aikido_zen.sinks.socket
 import aikido_zen.sinks.http_client
 import requests
+import urllib3
 
 
 @pytest.fixture(autouse=True)
@@ -19,7 +21,7 @@ def run_around_tests():
     get_cache().reset()
 
 
-def set_context_and_lifecycle(url):
+def set_context_and_lifecycle(url, host=None):
     wsgi_request = {
         "REQUEST_METHOD": "GET",
         "HTTP_HEADER_1": "header 1 value",
@@ -33,6 +35,8 @@ def set_context_and_lifecycle(url):
         "CONTENT_TYPE": "application/json",
         "REMOTE_ADDR": "198.51.100.23",
     }
+    if host is not None:
+        wsgi_request["HTTP_HOST"] = host
     context = Context(
         req=wsgi_request,
         body={
@@ -49,6 +53,9 @@ def ssrf_check(monkeypatch, url):
     monkeypatch.setenv("AIKIDO_BLOCK", "1")
     with pytest.raises(AikidoSSRF):
         requests.get(url)
+    with pytest.raises(AikidoSSRF):
+        http = urllib3.PoolManager()
+        http.request("GET", url)
 
 
 @pytest.mark.parametrize(
@@ -102,6 +109,18 @@ def test_no_raises_if_diff_url(monkeypatch):
         requests.get("http://ssrf-redirects.testssandbox.com/ssrf-test-domain-twice")
 
 
+def test_no_raises_if_diff_url_urllib3(monkeypatch):
+    http = urllib3.PoolManager()
+    set_context_and_lifecycle(
+        "http://firewallssrfredirects-env-2.eba-7ifve22q.eu-north-1.elasticbeanstalk.com/ssrf-test-domain-twice"
+    )
+    monkeypatch.setenv("AIKIDO_BLOCK", "1")
+    with pytest.raises(urllib3.exceptions.MaxRetryError):
+        http.request(
+            "GET", "http://ssrf-redirects.testssandbox.com/ssrf-test-domain-twice"
+        )
+
+
 def test_localhost_is_same_as_context(monkeypatch):
     set_context_and_lifecycle("http://localhost:8080")
     monkeypatch.setenv("AIKIDO_BLOCK", "1")
@@ -138,3 +157,31 @@ def test_different_capitalization_raises_ssrf(monkeypatch):
         requests.get("http://Localhost:8081/")
     with pytest.raises(AikidoSSRF):
         requests.get("http://localHost:8081/test")
+
+
+def test_srrf_with_request_to_itself_urllib3(monkeypatch):
+    http = urllib3.PoolManager()
+    reset_comms()
+
+    set_context_and_lifecycle("http://localhost:5000/test/1", host="localhost:5000")
+    monkeypatch.setenv("AIKIDO_BLOCK", "1")
+    with pytest.raises(urllib3.exceptions.MaxRetryError):
+        http.request("GET", "http://localhost:5000/test/1")
+
+    # Now test with no port match
+    set_context_and_lifecycle("http://localhost:5000/test/2", host="localhost:4999")
+    monkeypatch.setenv("AIKIDO_BLOCK", "1")
+    with pytest.raises(AikidoSSRF):
+        http.request("GET", "http://localhost:5000/test/2")
+
+    # Test with http app requesting to https
+    set_context_and_lifecycle("https://localhost/test/3", host="localhost:80")
+    monkeypatch.setenv("AIKIDO_BLOCK", "1")
+    with pytest.raises(urllib3.exceptions.MaxRetryError):
+        http.request("GET", "https://localhost/test/3")
+
+    # Test with https app requesting to http
+    set_context_and_lifecycle("http://localhost/test/4", host="localhost:443")
+    monkeypatch.setenv("AIKIDO_BLOCK", "1")
+    with pytest.raises(urllib3.exceptions.MaxRetryError):
+        http.request("GET", "https://localhost/test/4")