diff options
Diffstat (limited to 'src/hydrilla/proxy/addon.py')
| -rw-r--r-- | src/hydrilla/proxy/addon.py | 142 |
1 files changed, 84 insertions, 58 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) |