diff options
Diffstat (limited to 'src/hydrilla/proxy/policies')
-rw-r--r-- | src/hydrilla/proxy/policies/payload.py | 45 | ||||
-rw-r--r-- | src/hydrilla/proxy/policies/payload_resource.py | 2 |
2 files changed, 44 insertions, 3 deletions
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): diff --git a/src/hydrilla/proxy/policies/payload_resource.py b/src/hydrilla/proxy/policies/payload_resource.py index 3aacea0..f50f574 100644 --- a/src/hydrilla/proxy/policies/payload_resource.py +++ b/src/hydrilla/proxy/policies/payload_resource.py @@ -77,7 +77,7 @@ class PayloadResourcePolicy(PayloadAwarePolicy): -> http_messages.ProducedResponse: """....""" try: - file_data = self.payload_data.payload_ref.get_file_data(path) + file_data = self.payload_data.ref.get_file_data(path) except state.MissingItemError: return resource_blocked_response |