unfinished partial work

1 files changed, 171 insertions, 68 deletions

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)