diff --git a/bubble-server/src/main/resources/packer/roles/mitmproxy/files/bubble_api.py b/bubble-server/src/main/resources/packer/roles/mitmproxy/files/bubble_api.py index 76809dde..fc7320eb 100644 --- a/bubble-server/src/main/resources/packer/roles/mitmproxy/files/bubble_api.py +++ b/bubble-server/src/main/resources/packer/roles/mitmproxy/files/bubble_api.py @@ -450,7 +450,7 @@ def original_flex_ip(client_addr, fqdns): return None -def _replace_in_headers(headers, pattern: str, replacement: str): +def _replace_in_headers(headers: nheaders.Headers, pattern: str, replacement: str): """ Taken from original mitmproxy's Header class implementation with sligh change to allow replacement with empty string (resulting with actual removal/skip of the header line). @@ -489,17 +489,19 @@ def _replace_in_headers(headers, pattern: str, replacement: str): def response_header_modify(flow): - return None if flow.response is None else _header_modify(flow.client_conn.address[0], flow.response.headers) + if flow.response is None: + return None + return _header_modify(flow.client_conn.address[0], flow.server_conn.address[0], flow.response.headers) -def _header_modify(client_addr, headers): +def _header_modify(client_addr: str, server_addr: str, headers: nheaders.Headers): modifiers_set = 'responseHeaderModifierLists~' + client_addr + '~UNION' modifiers = REDIS.smembers(modifiers_set) repl_count = 0 if modifiers: for modifier in modifiers: - modifier_config = json.loads(modifier) + modifier_config = _extract_modifier_config(modifier, server_addr) repl_count += _replace_in_headers(headers, modifier_config['regex'], modifier_config['replacement']) if bubble_log.isEnabledFor(DEBUG): @@ -508,6 +510,11 @@ def _header_modify(client_addr, headers): return repl_count +def _extract_modifier_config(modifier: str, server_addr: str): + modifier.replace('{{fqdn}}', re.escape(server_addr)) + return json.loads(modifier) + + def health_check_response(flow): # if bubble_log.isEnabledFor(DEBUG): # bubble_log.debug('health_check_response: special bubble health check request, responding with OK')