From 9a3803d642e9488e4f39c3f57acf888b17dc71cd 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 | 25 ++++++++++++++++--- 1 file changed, 22 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..808a5cebaa611b 100644 --- a/test/integration/targets/setup_proxy/files/hamsandwich.py +++ b/test/integration/targets/setup_proxy/files/hamsandwich.py @@ -1,12 +1,31 @@ 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: + 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 self.parser.build_response()