diff options
Diffstat (limited to 'src/hydrilla/proxy/addon.py')
| -rw-r--r-- | src/hydrilla/proxy/addon.py | 93 |
1 files changed, 34 insertions, 59 deletions
diff --git a/src/hydrilla/proxy/addon.py b/src/hydrilla/proxy/addon.py index ae03ecc..de864fc 100644 --- a/src/hydrilla/proxy/addon.py +++ b/src/hydrilla/proxy/addon.py @@ -46,49 +46,25 @@ from mitmproxy.script import concurrent from ..exceptions import HaketiloException from ..translations import smart_gettext as _ -from ..url_patterns import parse_url, ParsedUrl +from .. import url_patterns from .state_impl import ConcreteHaketiloState from . import state from . import policies from . import http_messages -DefaultGetValue = t.TypeVar('DefaultGetValue', str, None) - -class MitmproxyHeadersWrapper(): - """....""" - def __init__(self, headers: http.Headers) -> None: - """....""" - self.headers = headers - - __getitem__ = lambda self, key: self.headers[key] - get_all = lambda self, key: self.headers.get_all(key) - - @t.overload - def get(self, key: str) -> t.Optional[str]: - ... - @t.overload - def get(self, key: str, default: DefaultGetValue) \ - -> t.Union[str, DefaultGetValue]: - ... - def get(self, key, default = None): - value = self.headers.get(key) - - if value is None: - return default - else: - return t.cast(str, value) - - def items(self) -> t.Iterable[tuple[str, str]]: - """....""" - return self.headers.items(multi=True) - - class LoggerToMitmproxy(state.Logger): def warn(self, msg: str) -> None: ctx.log.warn(f'Haketilo: {msg}') +def safe_parse_url(url: str) -> url_patterns.ParsedUrl: + try: + return url_patterns.parse_url(url) + except url_patterns.HaketiloURLException: + return url_patterns.dummy_url + + @dc.dataclass class FlowHandling: flow: http.HTTPFlow @@ -114,11 +90,10 @@ class FlowHandling: 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), + self._bl_response_info = http_messages.BodylessResponseInfo.make( + url = safe_parse_url(self.flow.request.url), status_code = self.flow.response.status_code, - headers = MitmproxyHeadersWrapper(headers) + headers = self.flow.response.headers ) return self._bl_response_info @@ -131,12 +106,15 @@ class FlowHandling: 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( + def make( + flow: http.HTTPFlow, + policy: policies.Policy, + url: url_patterns.ParsedUrl + ) -> 'FlowHandling': + bl_request_info = http_messages.BodylessRequestInfo.make( url = url, method = flow.request.method, - headers = MitmproxyHeadersWrapper(flow.request.headers) + headers = flow.request.headers ) return FlowHandling(flow, policy, bl_request_info) @@ -157,10 +135,6 @@ class PassedOptions: self.haketilo_launch_browser is not None) -magical_mitm_it_url_reg = re.compile(r'^http://mitm.it(/.*)?$') -dummy_url = parse_url('http://dummy.replacement.url') - - @dc.dataclass class HaketiloAddon: initial_options: PassedOptions = PassedOptions() @@ -257,16 +231,13 @@ class HaketiloAddon: handling = self.handling_dict.get(id(flow)) if handling is None: - parsed_url = dummy_url - - if magical_mitm_it_url_reg.match(flow.request.url): - policy = policies.DoNothingPolicy() + try: + parsed_url = url_patterns.parse_url(flow.request.url) + except url_patterns.HaketiloURLException as e: + policy = policies.ErrorBlockPolicy(builtin=True, error=e) + parsed_url = url_patterns.dummy_url else: - try: - parsed_url = parse_url(flow.request.url) - policy = self.state.select_policy(parsed_url) - except HaketiloException as e: - policy = policies.ErrorBlockPolicy(builtin=True, error=e) + policy = self.state.select_policy(parsed_url) handling = FlowHandling.make(flow, policy, parsed_url) @@ -330,16 +301,18 @@ class HaketiloAddon: result = handling.policy.consume_request(handling.request_info) if result is not None: - if isinstance(result, http_messages.ProducedRequest): - flow.request.url = result.url + mitmproxy_headers = http.Headers(result.headers.items_bin()) + + if isinstance(result, http_messages.RequestInfo): + flow.request.url = result.url.orig_url flow.request.method = result.method - flow.request.headers = http.Headers(result.headers) + flow.request.headers = mitmproxy_headers flow.request.set_content(result.body or None) else: - # isinstance(result, http_messages.ProducedResponse) + # isinstance(result, http_messages.ResponseInfo) flow.response = http.Response.make( status_code = result.status_code, - headers = http.Headers(result.headers), + headers = mitmproxy_headers, content = result.body ) @@ -370,8 +343,10 @@ class HaketiloAddon: response_info = handling.response_info ) if result is not None: + headers_bin = result.headers.items_bin() + flow.response.status_code = result.status_code - flow.response.headers = http.Headers(result.headers) + flow.response.headers = http.Headers(headers_bin) flow.response.set_content(result.body) self.forget_flow_handling(flow) |