From 2447f66eda99ee4567bff425d7f2fafeb8e9529b Mon Sep 17 00:00:00 2001 From: Matt Martz Date: Mon, 22 Apr 2024 15:11:27 -0500 Subject: [PATCH] improve plugin --- .../targets/setup_proxy/files/hamsandwich.py | 26 ++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/test/integration/targets/setup_proxy/files/hamsandwich.py b/test/integration/targets/setup_proxy/files/hamsandwich.py index ed7db8a4c6bb03..d3c487b34cefc8 100644 --- a/test/integration/targets/setup_proxy/files/hamsandwich.py +++ b/test/integration/targets/setup_proxy/files/hamsandwich.py @@ -1,12 +1,32 @@ from __future__ import annotations from proxy.http.proxy import HttpProxyBasePlugin +from proxy.http.parser import HttpParser, httpParserStates class HamSandwichPlugin(HttpProxyBasePlugin): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.parser = None + self.parsable = True + def handle_upstream_chunk(self, chunk): - headers, sep, body = chunk.tobytes().partition(b'\r\n\r\n') - if not sep: + if not self.parsable: + return chunk + + if not self.parser: + self.parser = HttpParser.response(chunk) + else: + self.parser.parse(chunk) + + if self.parser.state == httpParserStates.INITIALIZED: + # This is likely TLS without interception + self.parsable = False return chunk - return memoryview(bytearray(headers + b'\r\nX-Sandwich: ham' + sep + body)) + + if not self.parser.is_complete: + return None + + self.parser.add_header(b'X-Sandwich', b'ham') + return memoryview(bytearray(self.parser.build_response()))