diff options
Diffstat (limited to 'src/hydrilla')
| -rw-r--r-- | src/hydrilla/proxy/addon.py | 142 | ||||
| -rw-r--r-- | src/hydrilla/proxy/http_messages.py | 24 | ||||
| -rw-r--r-- | src/hydrilla/proxy/policies/base.py | 18 | ||||
| -rw-r--r-- | src/hydrilla/proxy/policies/misc.py | 7 | ||||
| -rw-r--r-- | src/hydrilla/proxy/policies/payload.py | 17 | ||||
| -rw-r--r-- | src/hydrilla/proxy/policies/payload_resource.py | 22 | ||||
| -rw-r--r-- | src/hydrilla/proxy/policies/rule.py | 8 |
7 files changed, 148 insertions, 90 deletions
diff --git a/src/hydrilla/proxy/addon.py b/src/hydrilla/proxy/addon.py index cca7924..ae03ecc 100644 --- a/src/hydrilla/proxy/addon.py +++ b/src/hydrilla/proxy/addon.py @@ -89,10 +89,57 @@ class LoggerToMitmproxy(state.Logger): ctx.log.warn(f'Haketilo: {msg}') -@dc.dataclass(frozen=True) -class FlowHandlingData: - request_url: ParsedUrl - policy: policies.Policy +@dc.dataclass +class FlowHandling: + flow: http.HTTPFlow + policy: policies.Policy + _bl_request_info: http_messages.BodylessRequestInfo + _request_info: t.Optional[http_messages.RequestInfo] = None + _bl_response_info: t.Optional[http_messages.BodylessResponseInfo] = None + + @property + def bl_request_info(self) -> http_messages.BodylessRequestInfo: + return self._bl_request_info + + @property + def request_info(self) -> http_messages.RequestInfo: + if self._request_info is None: + body = self.flow.request.get_content(strict=False) or b'' + self._request_info = self._bl_request_info.with_body(body) + + return self._request_info + + @property + def bl_response_info(self) -> http_messages.BodylessResponseInfo: + if self._bl_response_info is None: + assert self.flow.response is not None + + headers = self.flow.response.headers + self._bl_response_info = http_messages.BodylessResponseInfo( + url = parse_url(self.flow.request.url), + status_code = self.flow.response.status_code, + headers = MitmproxyHeadersWrapper(headers) + ) + + return self._bl_response_info + + @property + def response_info(self) -> http_messages.ResponseInfo: + assert self.flow.response is not None + + body = self.flow.response.get_content(strict=False) or b'' + return self.bl_response_info.with_body(body) + + @staticmethod + def make(flow: http.HTTPFlow, policy: policies.Policy, url: ParsedUrl) \ + -> 'FlowHandling': + bl_request_info = http_messages.BodylessRequestInfo( + url = url, + method = flow.request.method, + headers = MitmproxyHeadersWrapper(flow.request.headers) + ) + + return FlowHandling(flow, policy, bl_request_info) @dc.dataclass @@ -120,8 +167,8 @@ class HaketiloAddon: configured: bool = False configured_lock: Lock = dc.field(default_factory=Lock) - flows_data: dict[int, FlowHandlingData] = dc.field(default_factory=dict) - flows_data_lock: Lock = dc.field(default_factory=Lock) + handling_dict: dict[int, FlowHandling] = dc.field(default_factory=dict) + handling_dict_lock: Lock = dc.field(default_factory=Lock) logger: LoggerToMitmproxy = dc.field(default_factory=LoggerToMitmproxy) @@ -201,15 +248,15 @@ class HaketiloAddon: if not self.state.launch_browser(): self.logger.warn(_('warn.proxy.couldnt_launch_browser')) - def get_handling_data(self, flow: http.HTTPFlow) -> FlowHandlingData: + def get_flow_handling(self, flow: http.HTTPFlow) -> FlowHandling: policy: policies.Policy assert self.state is not None - with self.flows_data_lock: - handling_data = self.flows_data.get(id(flow)) + with self.handling_dict_lock: + handling = self.handling_dict.get(id(flow)) - if handling_data is None: + if handling is None: parsed_url = dummy_url if magical_mitm_it_url_reg.match(flow.request.url): @@ -221,17 +268,16 @@ class HaketiloAddon: except HaketiloException as e: policy = policies.ErrorBlockPolicy(builtin=True, error=e) - handling_data = FlowHandlingData(parsed_url, policy) + handling = FlowHandling.make(flow, policy, parsed_url) - with self.flows_data_lock: - self.flows_data[id(flow)] = handling_data + with self.handling_dict_lock: + self.handling_dict[id(flow)] = handling - return handling_data + return handling - def forget_handling_data(self, flow: http.HTTPFlow) -> None: - """....""" - with self.flows_data_lock: - self.flows_data.pop(id(flow), None) + def forget_flow_handling(self, flow: http.HTTPFlow) -> None: + with self.handling_dict_lock: + self.handling_dict.pop(id(flow), None) @contextmanager def http_safe_event_handling(self, flow: http.HTTPFlow) -> t.Iterator: @@ -252,7 +298,7 @@ class HaketiloAddon: headers = [(b'Content-Type', b'text/plain; charset=utf-8')] ) - self.forget_handling_data(flow) + self.forget_flow_handling(flow) @concurrent def requestheaders(self, flow: http.HTTPFlow) -> None: @@ -265,33 +311,23 @@ class HaketiloAddon: # visited before. flow.request.headers.pop('referer', None) - handling_data = self.get_handling_data(flow) - policy = handling_data.policy + handling = self.get_flow_handling(flow) + policy = handling.policy - if not policy.should_process_request(handling_data.request_url): + if not policy.should_process_request(handling.bl_request_info): flow.request.stream = True if policy.anticache: flow.request.anticache() @concurrent def request(self, flow: http.HTTPFlow) -> None: - """ - .... - """ if flow.request.stream: return with self.http_safe_event_handling(flow): - handling_data = self.get_handling_data(flow) + handling = self.get_flow_handling(flow) - request_info = http_messages.RequestInfo( - url = handling_data.request_url, - method = flow.request.method, - headers = MitmproxyHeadersWrapper(flow.request.headers), - body = flow.request.get_content(strict=False) or b'' - ) - - result = handling_data.policy.consume_request(request_info) + result = handling.policy.consume_request(handling.request_info) if result is not None: if isinstance(result, http_messages.ProducedRequest): @@ -303,51 +339,42 @@ class HaketiloAddon: # isinstance(result, http_messages.ProducedResponse) flow.response = http.Response.make( status_code = result.status_code, - headers = http.Headers(result.headers), - content = result.body + headers = http.Headers(result.headers), + content = result.body ) def responseheaders(self, flow: http.HTTPFlow) -> None: - """ - ...... - """ assert flow.response is not None with self.http_safe_event_handling(flow): - handling_data = self.get_handling_data(flow) - policy = handling_data.policy + handling = self.get_flow_handling(flow) - if not policy.should_process_response(handling_data.request_url): + if not handling.policy.should_process_response( + request_info = handling.request_info, + response_info = handling.bl_response_info + ): flow.response.stream = True @concurrent def response(self, flow: http.HTTPFlow) -> None: - """ - ...... - """ assert flow.response is not None if flow.response.stream: return with self.http_safe_event_handling(flow): - handling_data = self.get_handling_data(flow) - - response_info = http_messages.ResponseInfo( - url = parse_url(flow.request.url), - orig_url = handling_data.request_url, - status_code = flow.response.status_code, - headers = MitmproxyHeadersWrapper(flow.response.headers), - body = flow.response.get_content(strict=False) or b'' - ) + handling = self.get_flow_handling(flow) - result = handling_data.policy.consume_response(response_info) + result = handling.policy.consume_response( + request_info = handling.request_info, + response_info = handling.response_info + ) if result is not None: flow.response.status_code = result.status_code flow.response.headers = http.Headers(result.headers) flow.response.set_content(result.body) - self.forget_handling_data(flow) + self.forget_flow_handling(flow) def tls_clienthello(self, data: tls.ClientHelloData): if data.context.server.address is None: @@ -361,5 +388,4 @@ class HaketiloAddon: data.establish_server_tls_first = True def error(self, flow: http.HTTPFlow) -> None: - """....""" - self.forget_handling_data(flow) + self.forget_flow_handling(flow) 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()), diff --git a/src/hydrilla/proxy/policies/base.py b/src/hydrilla/proxy/policies/base.py index b3d3172..fcdbf9d 100644 --- a/src/hydrilla/proxy/policies/base.py +++ b/src/hydrilla/proxy/policies/base.py @@ -61,10 +61,17 @@ class Policy(ABC): priority: t.ClassVar[PolicyPriority] - def should_process_request(self, parsed_url: ParsedUrl) -> bool: + def should_process_request( + self, + request_info: http_messages.BodylessRequestInfo + ) -> bool: return self._process_request - def should_process_response(self, parsed_url: ParsedUrl) -> bool: + def should_process_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.BodylessResponseInfo + ) -> bool: return self._process_response def consume_request(self, request_info: http_messages.RequestInfo) \ @@ -73,8 +80,11 @@ class Policy(ABC): 'This kind of policy does not consume requests.' ) - def consume_response(self, response_info: http_messages.ResponseInfo) \ - -> t.Optional[http_messages.ProducedResponse]: + def consume_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.ResponseInfo + ) -> t.Optional[http_messages.ProducedResponse]: raise NotImplementedError( 'This kind of policy does not consume responses.' ) diff --git a/src/hydrilla/proxy/policies/misc.py b/src/hydrilla/proxy/policies/misc.py index 6d1e032..71692b3 100644 --- a/src/hydrilla/proxy/policies/misc.py +++ b/src/hydrilla/proxy/policies/misc.py @@ -71,8 +71,11 @@ class DoNothingPolicy(base.Policy): def consume_request(self, request_info: http_messages.RequestInfo) -> None: return None - def consume_response(self, response_info: http_messages.ResponseInfo) \ - -> None: + def consume_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.ResponseInfo + ) -> None: return None builtin: bool = True diff --git a/src/hydrilla/proxy/policies/payload.py b/src/hydrilla/proxy/policies/payload.py index d3e8e30..5b71af7 100644 --- a/src/hydrilla/proxy/policies/payload.py +++ b/src/hydrilla/proxy/policies/payload.py @@ -257,9 +257,11 @@ class PayloadInjectPolicy(PayloadAwarePolicy): return dc.replace(new_response, body=new_data) - def consume_response(self, response_info: http_messages.ResponseInfo) \ - -> http_messages.ProducedResponse: - """....""" + def consume_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.ResponseInfo + ) -> http_messages.ProducedResponse: try: return self._consume_response_unsafe(response_info) except Exception as e: @@ -286,15 +288,18 @@ class AutoPayloadInjectPolicy(PayloadInjectPolicy): """....""" priority: t.ClassVar[base.PolicyPriority] = base.PolicyPriority._ONE - def consume_response(self, response_info: http_messages.ResponseInfo) \ - -> http_messages.ProducedResponse: + def consume_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.ResponseInfo + ) -> http_messages.ProducedResponse: try: if self.payload_data.ref.has_problems(): raise _PayloadHasProblemsError() self.payload_data.ref.ensure_items_installed() - return super().consume_response(response_info) + return super().consume_response(request_info, response_info) except (state.RepoCommunicationError, state.FileInstallationError, _PayloadHasProblemsError) as ex: extra_params: dict[str, str] = { diff --git a/src/hydrilla/proxy/policies/payload_resource.py b/src/hydrilla/proxy/policies/payload_resource.py index d53f1f7..ae6a490 100644 --- a/src/hydrilla/proxy/policies/payload_resource.py +++ b/src/hydrilla/proxy/policies/payload_resource.py @@ -242,8 +242,12 @@ class PayloadResourcePolicy(PayloadAwarePolicy): segments_to_drop = len(self.payload_data.pattern_path_segments) + 1 return request_url.path_segments[segments_to_drop:] - def should_process_response(self, request_url: ParsedUrl) -> bool: - return self.extract_resource_path(request_url) \ + def should_process_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.BodylessResponseInfo + ) -> bool: + return self.extract_resource_path(request_info.url) \ == ('api', 'unrestricted_http') def _make_file_resource_response(self, path: tuple[str, ...]) \ @@ -338,8 +342,11 @@ class PayloadResourcePolicy(PayloadAwarePolicy): else: return resource_blocked_response - def consume_response(self, response_info: http_messages.ResponseInfo) \ - -> http_messages.ProducedResponse: + def consume_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.ResponseInfo + ) -> http_messages.ProducedResponse: """ This method shall only be called for responses to unrestricted HTTP API requests. Its purpose is to sanitize response headers and smuggle their @@ -351,7 +358,7 @@ class PayloadResourcePolicy(PayloadAwarePolicy): if (300 <= response_info.status_code < 400): location = response_info.headers.get('location') if location is not None: - orig_params = parse_qs(response_info.orig_url.query) + orig_params = parse_qs(request_info.url.query) orig_extra_headers_str, = orig_params['extra_headers'] new_query = urlencode({ @@ -359,10 +366,7 @@ class PayloadResourcePolicy(PayloadAwarePolicy): 'extra_headers': orig_extra_headers_str }) - new_url = urljoin( - response_info.orig_url.orig_url, - '?' + new_query - ) + new_url = urljoin(request_info.url.orig_url, '?' + new_query) extra_headers.append(('location', new_url)) diff --git a/src/hydrilla/proxy/policies/rule.py b/src/hydrilla/proxy/policies/rule.py index b742a64..2e9443e 100644 --- a/src/hydrilla/proxy/policies/rule.py +++ b/src/hydrilla/proxy/policies/rule.py @@ -79,9 +79,11 @@ class BlockPolicy(base.Policy): yield b'Content-Security-Policy', extra_csp.encode() - def consume_response(self, response_info: http_messages.ResponseInfo) \ - -> http_messages.ProducedResponse: - """....""" + def consume_response( + self, + request_info: http_messages.RequestInfo, + response_info: http_messages.ResponseInfo + ) -> http_messages.ProducedResponse: new_response = response_info.make_produced_response() new_headers = self._modify_headers(response_info) |