Linting & clarification

Author: bitterpanda63

aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py

@@ -108,7 +108,9 @@ def test_update_service_config_failure():
 
     # Set initial values
     connection_manager.conf.set_block_new_outgoing_requests(True)
-    connection_manager.conf.update_outbound_domains([{"hostname": "test.com", "mode": "block"}])
+    connection_manager.conf.update_outbound_domains(
+        [{"hostname": "test.com", "mode": "block"}]
+    )
 
     # Test failed response
     res = {"success": False, "blockNewOutgoingRequests": False, "domains": []}

aikido_zen/background_process/service_config.py

@@ -78,7 +78,9 @@ def is_bypassed_ip(self, ip):
         return self.bypassed_ips.has(ip)
 
     def update_outbound_domains(self, domains):
-        self.outbound_domains = {domain["hostname"]: domain["mode"] for domain in domains}
+        self.outbound_domains = {
+            domain["hostname"]: domain["mode"] for domain in domains
+        }
 
     def set_block_new_outgoing_requests(self, value: bool):
         """Set whether to block new outgoing requests"""

aikido_zen/sinks/socket/__init__.py

@@ -1,11 +1,14 @@
 """
 Sink module for `socket`
 """
+
 from aikido_zen.errors import AikidoSSRF
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.sinks import on_import, patch_function, after
 from aikido_zen.sinks.socket.normalize_hostname import normalize_hostname
-from aikido_zen.sinks.socket.should_block_outbound_domain import should_block_outbound_domain
+from aikido_zen.sinks.socket.should_block_outbound_domain import (
+    should_block_outbound_domain,
+)
 from aikido_zen.vulnerabilities import run_vulnerability_scan
 
 

aikido_zen/sinks/socket/should_block_outbound_domain.py

@@ -1,13 +1,14 @@
-from aikido_zen.context import get_current_context
-from aikido_zen.errors import AikidoSSRF
-from aikido_zen.sinks.socket.normalize_hostname import normalize_hostname
 from aikido_zen.thread.thread_cache import get_cache
 
 
 def should_block_outbound_domain(hostname, port):
     process_cache = get_cache()
     if not process_cache:
         return False
+
+    # We store the hostname before checking the blocking status
+    # This is because if we are in lockdown mode and blocking all new hostnames, it should still
+    # show up in the dashboard. This allows the user to allow traffic to newly detected hostnames.
     process_cache.hostnames.add(hostname, port)
 
     return process_cache.config.should_block_outgoing_request(hostname)