aboutsummaryrefslogtreecommitdiff
path: root/src/hydrilla/proxy/policies/payload.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/hydrilla/proxy/policies/payload.py')
-rw-r--r--src/hydrilla/proxy/policies/payload.py45
1 files changed, 43 insertions, 2 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):