From 3beab050c339c51c484af9bcd8248ba8ebbbf4d4 Mon Sep 17 00:00:00 2001
From: Wojtek Kosior <koszko@koszko.org>
Date: Wed, 19 Oct 2022 18:55:08 +0200
Subject: [proxy] pass all available flow information to relevant methods of
 Policy type

---
 src/hydrilla/proxy/http_messages.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

(limited to 'src/hydrilla/proxy/http_messages.py')

diff --git a/src/hydrilla/proxy/http_messages.py b/src/hydrilla/proxy/http_messages.py
index 78baf81..dbf2c63 100644
--- a/src/hydrilla/proxy/http_messages.py
+++ b/src/hydrilla/proxy/http_messages.py
@@ -74,15 +74,19 @@ class ProducedRequest:
     body:        bytes
 
 @dc.dataclass(frozen=True)
-class RequestInfo:
-    """...."""
+class BodylessRequestInfo:
     url:         url_patterns.ParsedUrl
     method:      str
     headers:     IHeaders
-    body:        bytes
+
+    def with_body(self, body: bytes) -> 'RequestInfo':
+        return RequestInfo(self.url, self.method, self.headers, body)
+
+@dc.dataclass(frozen=True)
+class RequestInfo(BodylessRequestInfo):
+    body: bytes
 
     def make_produced_request(self) -> ProducedRequest:
-        """...."""
         return ProducedRequest(
             url     = self.url.orig_url,
             method  = self.method,
@@ -98,16 +102,20 @@ class ProducedResponse:
     body:        bytes
 
 @dc.dataclass(frozen=True)
-class ResponseInfo:
+class BodylessResponseInfo:
     """...."""
     url:         url_patterns.ParsedUrl
-    orig_url:    url_patterns.ParsedUrl
     status_code: int
     headers:     IHeaders
-    body:        bytes
+
+    def with_body(self, body: bytes) -> 'ResponseInfo':
+        return ResponseInfo(self.url, self.status_code, self.headers, body)
+
+@dc.dataclass(frozen=True)
+class ResponseInfo(BodylessResponseInfo):
+    body: bytes
 
     def make_produced_response(self) -> ProducedResponse:
-        """...."""
         return ProducedResponse(
             status_code = self.status_code,
             headers     = encode_headers_items(self.headers.items()),
-- 
cgit v1.2.3