From 45e5cf8dc3ca936e2db8e7e45689d0a3331aad43 Mon Sep 17 00:00:00 2001 From: Wojtek Kosior Date: Thu, 8 Sep 2022 17:59:40 +0200 Subject: [proxy] make package auto-installation work --- src/hydrilla/proxy/policies/payload.py | 45 ++++++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 2 deletions(-) (limited to 'src/hydrilla/proxy/policies/payload.py') diff --git a/src/hydrilla/proxy/policies/payload.py b/src/hydrilla/proxy/policies/payload.py index e64653d..a063c7c 100644 --- a/src/hydrilla/proxy/policies/payload.py +++ b/src/hydrilla/proxy/policies/payload.py @@ -35,8 +35,12 @@ import dataclasses as dc import typing as t import re +from urllib.parse import urlencode + +from itsdangerous.url_safe import URLSafeSerializer import bs4 # type: ignore +from ...exceptions import HaketiloException from ...url_patterns import ParsedUrl from .. import csp from .. import state @@ -65,7 +69,7 @@ class PayloadAwarePolicyFactory(base.PolicyFactory): @property def payload_ref(self) -> state.PayloadRef: """....""" - return self.payload_key.payload_ref + return self.payload_key.ref def __lt__(self, other: base.PolicyFactory) -> bool: """....""" @@ -171,7 +175,7 @@ class PayloadInjectPolicy(PayloadAwarePolicy): def _script_urls(self, url: ParsedUrl) -> t.Iterable[str]: """....""" base_url = self.assets_base_url(url) - payload_ref = self.payload_data.payload_ref + payload_ref = self.payload_data.ref for path in payload_ref.get_script_paths(): yield base_url + '/'.join(('static', *path)) @@ -261,10 +265,47 @@ class PayloadInjectPolicy(PayloadAwarePolicy): ) +class _PayloadHasProblemsError(HaketiloException): + pass + class AutoPayloadInjectPolicy(PayloadInjectPolicy): """....""" priority: t.ClassVar[base.PolicyPriority] = base.PolicyPriority._ONE + def consume_response(self, 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) + except (state.RepoCommunicationError, state.FileInstallationError, + _PayloadHasProblemsError) as ex: + params = { + 'next_url': response_info.url.orig_url, + 'payload_id': self.payload_data.ref.id + } + + if isinstance(ex, state.FileInstallationError): + params['repo_id'] = ex.repo_id + params['file_sha256'] = ex.sha256 + + serializer = URLSafeSerializer( + self.payload_data.global_secret, + salt = 'auto_install_error' + ) + query = urlencode({'details': serializer.dumps(params)}) + redirect_url = 'https://hkt.mitm.it/auto_install_error?' + query + msg = 'Error occured when installing payload. Redirecting.' + + return http_messages.ProducedResponse( + status_code = 303, + headers = [(b'Location', redirect_url.encode())], + body = msg.encode() + ) + @dc.dataclass(frozen=True) class PayloadSuggestPolicy(PayloadAwarePolicy): -- cgit v1.2.3