diff options
Diffstat (limited to 'src/hydrilla/proxy/policies/rule.py')
| -rw-r--r-- | src/hydrilla/proxy/policies/rule.py | 23 |
1 files changed, 12 insertions, 11 deletions
diff --git a/src/hydrilla/proxy/policies/rule.py b/src/hydrilla/proxy/policies/rule.py index 2e9443e..8272d2f 100644 --- a/src/hydrilla/proxy/policies/rule.py +++ b/src/hydrilla/proxy/policies/rule.py @@ -50,13 +50,14 @@ class BlockPolicy(base.Policy): priority: t.ClassVar[base.PolicyPriority] = base.PolicyPriority._TWO def _modify_headers(self, response_info: http_messages.ResponseInfo) \ - -> t.Iterable[tuple[bytes, bytes]]: - """....""" + -> http_messages.IHeaders: + new_headers = [] + csp_policies = csp.extract(response_info.headers) - for header_name, header_value in response_info.headers.items(): - if header_name.lower() not in csp.header_names_and_dispositions: - yield header_name.encode(), header_value.encode() + for key, val in response_info.headers.items(): + if key.lower() not in csp.header_names_and_dispositions: + new_headers.append((key, val)) for policy in csp_policies: if policy.disposition != 'enforce': @@ -68,7 +69,7 @@ class BlockPolicy(base.Policy): policy = dc.replace(policy, directives=directives.finish()) - yield policy.header_name.encode(), policy.serialize().encode() + new_headers.append((policy.header_name, policy.serialize())) extra_csp = ';'.join(( "script-src 'none'", @@ -76,19 +77,19 @@ class BlockPolicy(base.Policy): "script-src-attr 'none'" )) - yield b'Content-Security-Policy', extra_csp.encode() + new_headers.append(('Content-Security-Policy', extra_csp)) + + return http_messages.make_headers(new_headers) def consume_response( self, request_info: http_messages.RequestInfo, response_info: http_messages.ResponseInfo - ) -> http_messages.ProducedResponse: - new_response = response_info.make_produced_response() - + ) -> http_messages.ResponseInfo: new_headers = self._modify_headers(response_info) - return dc.replace(new_response, headers=new_headers) + return dc.replace(response_info, headers=new_headers) @dc.dataclass(frozen=True) class RuleAllowPolicy(AllowPolicy): |