diff options
Diffstat (limited to 'src/hydrilla/proxy/policies/base.py')
| -rw-r--r-- | src/hydrilla/proxy/policies/base.py | 69 |
1 files changed, 62 insertions, 7 deletions
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): """....""" |