[proxy] make package auto-installation work

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):