From 879c41927171efc8d77d1de2739b18e2eb57580f Mon Sep 17 00:00:00 2001 From: Wojtek Kosior Date: Wed, 27 Jul 2022 15:56:24 +0200 Subject: unfinished partial work --- src/hydrilla/proxy/addon.py | 239 +++++++++++++++++++++++++++++++------------- 1 file changed, 171 insertions(+), 68 deletions(-) (limited to 'src/hydrilla/proxy/addon.py') diff --git a/src/hydrilla/proxy/addon.py b/src/hydrilla/proxy/addon.py index 7d6487b..16c2841 100644 --- a/src/hydrilla/proxy/addon.py +++ b/src/hydrilla/proxy/addon.py @@ -32,58 +32,70 @@ from addon script. # Enable using with Python 3.7. from __future__ import annotations -import os.path +import sys import typing as t import dataclasses as dc +import traceback as tb from threading import Lock from pathlib import Path from contextlib import contextmanager -from mitmproxy import http, addonmanager, ctx +# for mitmproxy 6.* +from mitmproxy.net import http +if not hasattr(http, 'Headers'): + # for mitmproxy 8.* + from mitmproxy import http # type: ignore + +from mitmproxy.http import HTTPFlow + +from mitmproxy import addonmanager, ctx from mitmproxy.script import concurrent -from .flow_handlers import make_flow_handler, FlowHandler -from .state import HaketiloState +from ..exceptions import HaketiloException from ..translations import smart_gettext as _ +from ..url_patterns import parse_url +from .state_impl import ConcreteHaketiloState +from . import policies -FlowHandlers = dict[int, FlowHandler] -StateUpdater = t.Callable[[HaketiloState], None] +DefaultGetValue = t.TypeVar('DefaultGetValue', object, None) -HTTPHandlerFun = t.Callable[ - ['HaketiloAddon', http.HTTPFlow], - t.Optional[StateUpdater] -] - -def http_event_handler(handler_fun: HTTPHandlerFun): - """....decorator""" - def wrapped_handler(self: 'HaketiloAddon', flow: http.HTTPFlow): +class MitmproxyHeadersWrapper(): + """....""" + def __init__(self, headers: http.Headers) -> None: """....""" - with self.configured_lock: - assert self.configured + self.headers = headers - assert self.state is not None + __getitem__ = lambda self, key: self.headers[key] + get_all = lambda self, key: self.headers.get_all(key) - state_updater = handler_fun(self, flow) + def get(self, key: str, default: DefaultGetValue = None) \ + -> t.Union[str, DefaultGetValue]: + """....""" + value = self.headers.get(key) - if state_updater is not None: - state_updater(self.state) + if value is None: + return default + else: + return t.cast(str, value) - return wrapped_handler + def items(self) -> t.Iterable[tuple[str, str]]: + """....""" + return self.headers.items(multi=True) @dc.dataclass class HaketiloAddon: """ ....... """ - configured: bool = False - configured_lock: Lock = dc.field(default_factory=Lock) + configured: bool = False + configured_lock: Lock = dc.field(default_factory=Lock) - state: t.Optional[HaketiloState] = None + flow_policies: dict[int, policies.Policy] = dc.field(default_factory=dict) + policies_lock: Lock = dc.field(default_factory=Lock) - flow_handlers: FlowHandlers = dc.field(default_factory=dict) - handlers_lock: Lock = dc.field(default_factory=Lock) + state: t.Optional[ConcreteHaketiloState] = None def load(self, loader: addonmanager.Loader) -> None: """....""" @@ -104,74 +116,165 @@ class HaketiloAddon: ctx.log.warn(_('haketilo_dir_already_configured')) return - haketilo_dir = Path(ctx.options.haketilo_dir) - self.state = HaketiloState(haketilo_dir / 'store') + try: + haketilo_dir = Path(ctx.options.haketilo_dir) + + self.state = ConcreteHaketiloState.make(haketilo_dir / 'store') + except Exception as e: + tb.print_exception(None, e, e.__traceback__) + sys.exit(1) - def assign_handler(self, flow: http.HTTPFlow, flow_handler: FlowHandler) \ - -> None: + self.configured = True + + def try_get_policy(self, flow: HTTPFlow, fail_ok: bool = True) -> \ + t.Optional[policies.Policy]: """....""" - with self.handlers_lock: - self.flow_handlers[id(flow)] = flow_handler + with self.policies_lock: + policy = self.flow_policies.get(id(flow)) + + if policy is None: + try: + parsed_url = parse_url(flow.request.url) + except HaketiloException: + if fail_ok: + return None + else: + raise + + assert self.state is not None + + policy = self.state.select_policy(parsed_url) + + with self.policies_lock: + self.flow_policies[id(flow)] = policy + + return policy + + def get_policy(self, flow: HTTPFlow) -> policies.Policy: + return t.cast(policies.Policy, self.try_get_policy(flow, fail_ok=False)) - def lookup_handler(self, flow: http.HTTPFlow) -> FlowHandler: + def forget_policy(self, flow: HTTPFlow) -> None: """....""" - with self.handlers_lock: - return self.flow_handlers[id(flow)] + with self.policies_lock: + self.flow_policies.pop(id(flow), None) - def forget_handler(self, flow: http.HTTPFlow) -> None: + @contextmanager + def http_safe_event_handling(self, flow: HTTPFlow) -> t.Iterator: """....""" - with self.handlers_lock: - self.flow_handlers.pop(id(flow), None) + with self.configured_lock: + assert self.configured + + try: + yield + except Exception as e: + tb_string = ''.join(tb.format_exception(None, e, e.__traceback__)) + error_text = _('err.proxy.unknown_error_{}_try_again')\ + .format(tb_string)\ + .encode() + flow.response = http.Response.make( + status_code = 500, + content = error_text, + headers = [(b'Content-Type', b'text/plain; charset=utf-8')] + ) + + self.forget_policy(flow) @concurrent - @http_event_handler - def requestheaders(self, flow: http.HTTPFlow) -> t.Optional[StateUpdater]: + def requestheaders(self, flow: HTTPFlow) -> None: + # TODO: don't account for mitmproxy 6 in the code + # Mitmproxy 6 causes even more strange behavior than described below. + # This cannot be easily worked around. Let's just use version 8 and + # make an APT package for it. """ - ..... + Under mitmproxy 8 this handler deduces an appropriate policy for flow's + URL and assigns it to the flow. Under mitmproxy 6 the URL is not yet + available at this point, so the handler effectively does nothing. """ - assert self.state is not None - - policy = self.state.select_policy(flow.request.url) + with self.http_safe_event_handling(flow): + policy = self.try_get_policy(flow) - flow_handler = make_flow_handler(flow, policy) - - self.assign_handler(flow, flow_handler) - - return flow_handler.on_requestheaders() + if policy is not None: + if not policy.process_request: + flow.request.stream = True + if policy.anticache: + flow.request.anticache() @concurrent - @http_event_handler - def request(self, flow: http.HTTPFlow) -> t.Optional[StateUpdater]: + def request(self, flow: HTTPFlow) -> None: """ .... """ - return self.lookup_handler(flow).on_request() + if flow.request.stream: + return - @concurrent - @http_event_handler - def responseheaders(self, flow: http.HTTPFlow) -> t.Optional[StateUpdater]: + with self.http_safe_event_handling(flow): + policy = self.get_policy(flow) + + request_info = policies.RequestInfo( + url = parse_url(flow.request.url), + method = flow.request.method, + headers = MitmproxyHeadersWrapper(flow.request.headers), + body = flow.request.get_content(strict=False) or b'' + ) + + result = policy.consume_request(request_info) + + if result is not None: + if isinstance(result, policies.ProducedRequest): + flow.request = http.Request.make( + url = result.url, + method = result.method, + headers = http.Headers(result.headers), + content = result.body + ) + else: + # isinstance(result, policies.ProducedResponse) + flow.response = http.Response.make( + status_code = result.status_code, + headers = http.Headers(result.headers), + content = result.body + ) + + def responseheaders(self, flow: HTTPFlow) -> None: """ ...... """ - return self.lookup_handler(flow).on_responseheaders() + assert flow.response is not None + + with self.http_safe_event_handling(flow): + policy = self.get_policy(flow) + + if not policy.process_response: + flow.response.stream = True @concurrent - @http_event_handler - def response(self, flow: http.HTTPFlow) -> t.Optional[StateUpdater]: + def response(self, flow: HTTPFlow) -> None: """ ...... """ - updater = self.lookup_handler(flow).on_response() + assert flow.response is not None - self.forget_handler(flow) + if flow.response.stream: + return - return updater + with self.http_safe_event_handling(flow): + policy = self.get_policy(flow) - @http_event_handler - def error(self, flow: http.HTTPFlow) -> None: - """....""" - self.forget_handler(flow) + response_info = policies.ResponseInfo( + url = parse_url(flow.request.url), + status_code = flow.response.status_code, + headers = MitmproxyHeadersWrapper(flow.response.headers), + body = flow.response.get_content(strict=False) or b'' + ) -addons = [ - HaketiloAddon() -] + result = policy.consume_response(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_policy(flow) + + def error(self, flow: HTTPFlow) -> None: + """....""" + self.forget_policy(flow) -- cgit v1.2.3