[proxy] Add support for script blocking/allowing rules

1 files changed, 35 insertions, 10 deletions

diff --git a/src/hydrilla/proxy/state_impl/concrete_state.py b/src/hydrilla/proxy/state_impl/concrete_state.py
index cd32e83..a6d32f1 100644
--- a/src/hydrilla/proxy/state_impl/concrete_state.py
+++ b/src/hydrilla/proxy/state_impl/concrete_state.py
@@ -47,6 +47,7 @@ from .. import state as st
 from .. import policies
 from .. import simple_dependency_satisfying as sds
 from . import base
+from . import rules
 from . import items
 from . import repos
 from . import payloads
@@ -143,7 +144,7 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
                 repo_id         = repo_id
             )
 
-            self.rebuild_structures()
+            self.rebuild_structures(rules=False)
 
     def prune_orphans(self) -> None:
         with self.cursor() as cursor:
@@ -163,7 +164,7 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
                 unlocked_required_mappings = unlocked_required_mappings
             )
 
-            self.rebuild_structures()
+            self.rebuild_structures(rules=False)
 
     def pull_missing_files(self) -> None:
         with self.cursor() as cursor:
@@ -182,6 +183,29 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
                 ui_factory
             )
 
+        # Put script blocking/allowing rules in policy tree.
+        cursor.execute('SELECT pattern, allow_scripts FROM rules;')
+
+        for pattern, allow_scripts in cursor.fetchall():
+            for parsed_pattern in url_patterns.parse_pattern(pattern):
+                factory: policies.PolicyFactory
+                if allow_scripts:
+                    factory = policies.RuleAllowPolicyFactory(
+                        builtin = False,
+                        pattern = parsed_pattern
+                    )
+                else:
+                    factory = policies.RuleBlockPolicyFactory(
+                        builtin = False,
+                        pattern = parsed_pattern
+                    )
+
+                new_policy_tree = new_policy_tree.register(
+                    parsed_pattern = parsed_pattern,
+                    item           = factory
+                )
+
+        # Put script payload rules in policy tree.
         cursor.execute(
             '''
             SELECT
@@ -202,15 +226,10 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
             '''
         )
 
-        rows = cursor.fetchall()
-
         new_payloads_data: dict[st.PayloadRef, st.PayloadData] = {}
 
-        for row in rows:
-            (payload_id_int, pattern, eval_allowed, cors_bypass_allowed,
-             enabled_status,
-             identifier) = row
-
+        for (payload_id_int, pattern, eval_allowed, cors_bypass_allowed,
+             enabled_status, identifier) in cursor.fetchall():
             payload_ref = payloads.ConcretePayloadRef(str(payload_id_int), self)
 
             previous_data = self.payloads_data.get(payload_ref)
@@ -245,10 +264,16 @@ class ConcreteHaketiloState(base.HaketiloStateWithFields):
         self.policy_tree   = new_policy_tree
         self.payloads_data = new_payloads_data
 
-    def rebuild_structures(self) -> None:
+    def rebuild_structures(self, *, payloads: bool = True, rules: bool = True) \
+        -> None:
+        # The `payloads` and `rules` args will be useful for optimization but
+        # for now we're not yet using them.
         with self.cursor() as cursor:
             self._rebuild_structures(cursor)
 
+    def rule_store(self) -> st.RuleStore:
+        return rules.ConcreteRuleStore(self)
+
     def repo_store(self) -> st.RepoStore:
         return repos.ConcreteRepoStore(self)