2025-12-25 upload

venv/Lib/site-packages/mitmproxy/addons/stickycookie.py

@@ -0,0 +1,97 @@
+import collections
+from http import cookiejar
+from typing import Optional
+
+from mitmproxy import ctx
+from mitmproxy import exceptions
+from mitmproxy import flowfilter
+from mitmproxy import http
+from mitmproxy.net.http import cookies
+
+TOrigin = tuple[str, int, str]
+
+
+def ckey(attrs: dict[str, str], f: http.HTTPFlow) -> TOrigin:
+    """
+    Returns a (domain, port, path) tuple.
+    """
+    domain = f.request.host
+    path = "/"
+    if "domain" in attrs:
+        domain = attrs["domain"]
+    if "path" in attrs:
+        path = attrs["path"]
+    return (domain, f.request.port, path)
+
+
+def domain_match(a: str, b: str) -> bool:
+    if cookiejar.domain_match(a, b):  # type: ignore
+        return True
+    elif cookiejar.domain_match(a, b.strip(".")):  # type: ignore
+        return True
+    return False
+
+
+class StickyCookie:
+    def __init__(self) -> None:
+        self.jar: collections.defaultdict[TOrigin, dict[str, str]] = (
+            collections.defaultdict(dict)
+        )
+        self.flt: flowfilter.TFilter | None = None
+
+    def load(self, loader):
+        loader.add_option(
+            "stickycookie",
+            Optional[str],
+            None,
+            "Set sticky cookie filter. Matched against requests.",
+        )
+
+    def configure(self, updated):
+        if "stickycookie" in updated:
+            if ctx.options.stickycookie:
+                try:
+                    self.flt = flowfilter.parse(ctx.options.stickycookie)
+                except ValueError as e:
+                    raise exceptions.OptionsError(str(e)) from e
+            else:
+                self.flt = None
+
+    def response(self, flow: http.HTTPFlow):
+        assert flow.response
+        if self.flt:
+            for name, (value, attrs) in flow.response.cookies.items(multi=True):
+                # FIXME: We now know that Cookie.py screws up some cookies with
+                # valid RFC 822/1123 datetime specifications for expiry. Sigh.
+                dom_port_path = ckey(attrs, flow)
+
+                if domain_match(flow.request.host, dom_port_path[0]):
+                    if cookies.is_expired(attrs):
+                        # Remove the cookie from jar
+                        self.jar[dom_port_path].pop(name, None)
+
+                        # If all cookies of a dom_port_path have been removed
+                        # then remove it from the jar itself
+                        if not self.jar[dom_port_path]:
+                            self.jar.pop(dom_port_path, None)
+                    else:
+                        self.jar[dom_port_path][name] = value
+
+    def request(self, flow: http.HTTPFlow):
+        if self.flt:
+            cookie_list: list[tuple[str, str]] = []
+            if flowfilter.match(self.flt, flow):
+                for (domain, port, path), c in self.jar.items():
+                    match = [
+                        domain_match(flow.request.host, domain),
+                        flow.request.port == port,
+                        flow.request.path.startswith(path),
+                    ]
+                    if all(match):
+                        cookie_list.extend(c.items())
+            if cookie_list:
+                # FIXME: we need to formalise this...
+                flow.metadata["stickycookie"] = True
+                flow.request.headers["cookie"] = cookies.format_cookie_header(
+                    cookie_list
+                )