From fe43bd552aaacd649b0e00afada01d07ad8dae9a Mon Sep 17 00:00:00 2001
From: Wojtek Kosior <koszko@koszko.org>
Date: Fri, 21 Oct 2022 21:42:56 +0200
Subject: [proxy] facilitate injecting non-payload script to all pages

---
 src/hydrilla/proxy/policies/base.py | 69 +++++++++++++++++++++++++++++++++----
 1 file changed, 62 insertions(+), 7 deletions(-)

(limited to 'src/hydrilla/proxy/policies/base.py')

diff --git a/src/hydrilla/proxy/policies/base.py b/src/hydrilla/proxy/policies/base.py
index 7ce8663..0c37185 100644
--- a/src/hydrilla/proxy/policies/base.py
+++ b/src/hydrilla/proxy/policies/base.py
@@ -29,11 +29,17 @@
 .....
 """
 
+import enum
+import re
 import dataclasses as dc
 import typing as t
-import enum
 
+from threading import Lock
 from abc import ABC, abstractmethod
+from hashlib import sha256
+from base64 import b64encode
+
+import jinja2
 
 from immutables import Map
 
@@ -43,6 +49,19 @@ from .. import http_messages
 from .. import csp
 
 
+loader = jinja2.PackageLoader(__package__, package_path='injectable_scripts')
+jinja_env = jinja2.Environment(
+    loader        = loader,
+    lstrip_blocks = True,
+    autoescape    = False
+)
+jinja_lock = Lock()
+
+
+popup_script = jinja_env.get_template('popup.js.jinja').render()
+popup_script_sha256_bytes = sha256(popup_script.encode()).digest()
+popup_script_sha256_b64 = b64encode(popup_script_sha256_bytes).decode()
+
 class PolicyPriority(int, enum.Enum):
     """...."""
     _ONE   = 1
@@ -55,6 +74,11 @@ MessageInfo = t.Union[
 ]
 
 
+# We're doing *very* simple doctype matching for now. If a site wanted, it could
+# trick us into getting this wrong.
+doctype_re = re.compile(r'^\s*<!doctype[^>]*>', re.IGNORECASE)
+
+
 UTF8_BOM = b'\xEF\xBB\xBF'
 BOMs = (
     (UTF8_BOM,    'utf-8'),
@@ -63,14 +87,22 @@ BOMs = (
 )
 
 
+# mypy needs to be corrected:
+# https://stackoverflow.com/questions/70999513/conflict-between-mix-ins-for-abstract-dataclasses/70999704#70999704
+@dc.dataclass(frozen=True) # type: ignore[misc]
 class Policy(ABC):
-    """...."""
     _process_request:  t.ClassVar[bool] = False
     _process_response: t.ClassVar[bool] = False
     anticache:         t.ClassVar[bool] = True
 
     priority: t.ClassVar[PolicyPriority]
 
+    haketilo_settings: state.HaketiloGlobalSettings
+
+    @property
+    def current_popup_settings(self) -> state.PopupSettings:
+        return self.haketilo_settings.default_popup_jsallowed
+
     def should_process_request(
             self,
             request_info: http_messages.BodylessRequestInfo
@@ -82,7 +114,11 @@ class Policy(ABC):
             request_info:  http_messages.RequestInfo,
             response_info: http_messages.AnyResponseInfo
     ) -> bool:
-        return self._process_response
+        if self._process_response:
+            return True
+
+        return (self.current_popup_settings.popup_enabled and
+                http_messages.is_likely_a_page(request_info, response_info))
 
     def _csp_to_clear(self, http_info: http_messages.FullHTTPInfo) \
         -> t.Union[t.Sequence[str], t.Literal['all']]:
@@ -94,7 +130,11 @@ class Policy(ABC):
 
     def _csp_to_extend(self, http_info: http_messages.FullHTTPInfo) \
         -> t.Mapping[str, t.Sequence[str]]:
-        return Map()
+        if (self.current_popup_settings.popup_enabled and
+            http_info.is_likely_a_page):
+            return {'script-src': [f"'sha256-{popup_script_sha256_b64}'"]}
+        else:
+            return Map()
 
     def _modify_response_headers(self, http_info: http_messages.FullHTTPInfo) \
         -> http_messages.IHeaders:
@@ -117,7 +157,24 @@ class Policy(ABC):
             http_info: http_messages.FullHTTPInfo,
             encoding:  t.Optional[str]
     ) -> t.Union[str, bytes]:
-        return http_info.response_info.body
+        popup_settings = self.current_popup_settings
+
+        if (popup_settings.popup_enabled and
+            http_info.is_likely_a_page):
+            if encoding is None:
+                encoding = 'utf-8'
+
+            body_bytes = http_info.response_info.body
+            body = body_bytes.decode(encoding, errors='replace')
+
+            match = doctype_re.match(body)
+            doctype_decl_len = 0 if match is None else match.end()
+
+            dotype_decl = body[0:doctype_decl_len]
+            doc_rest = body[doctype_decl_len:]
+            return f'{dotype_decl}<script>{popup_script}</script>{doc_rest}'
+        else:
+            return http_info.response_info.body
 
     def _modify_response_body(self, http_info: http_messages.FullHTTPInfo) \
         -> bytes:
@@ -188,8 +245,6 @@ class Policy(ABC):
         )
 
 
-# mypy needs to be corrected:
-# https://stackoverflow.com/questions/70999513/conflict-between-mix-ins-for-abstract-dataclasses/70999704#70999704
 @dc.dataclass(frozen=True, unsafe_hash=True) # type: ignore[misc]
 class PolicyFactory(ABC):
     """...."""
-- 
cgit v1.2.3