2025-12-25 upload

venv/Lib/site-packages/mitmproxy/test/taddons.py

@@ -0,0 +1,82 @@
+import asyncio
+
+import mitmproxy.master
+import mitmproxy.options
+from mitmproxy import command
+from mitmproxy import eventsequence
+from mitmproxy import hooks
+from mitmproxy.addons import core
+from mitmproxy.addons import script
+
+
+class context:
+    """
+    A context for testing addons, which sets up the mitmproxy.ctx module so
+    handlers can run as they would within mitmproxy. The context also
+    provides a number of helper methods for common testing scenarios.
+    """
+
+    def __init__(self, *addons, options=None, loadcore=True):
+        self.owns_loop = False
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            self.owns_loop = True
+            loop = asyncio.new_event_loop()
+
+        options = options or mitmproxy.options.Options()
+        self.master = mitmproxy.master.Master(options, event_loop=loop)
+        self.options = self.master.options
+
+        if loadcore:
+            self.master.addons.add(core.Core())
+
+        for a in addons:
+            self.master.addons.add(a)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.owns_loop and not self.master.event_loop.is_closed():
+            # Close the loop if we created it
+            self.master.event_loop.close()
+        return False
+
+    async def cycle(self, addon, f):
+        """
+        Cycles the flow through the events for the flow. Stops if the flow
+        is intercepted.
+        """
+        for evt in eventsequence.iterate(f):
+            await self.master.addons.invoke_addon(addon, evt)
+            if f.intercepted:
+                return
+
+    def configure(self, addon, **kwargs):
+        """
+        A helper for testing configure methods. Modifies the registered
+        Options object with the given keyword arguments, then calls the
+        configure method on the addon with the updated value.
+        """
+        if addon not in self.master.addons:
+            self.master.addons.register(addon)
+        with self.options.rollback(kwargs.keys(), reraise=True):
+            if kwargs:
+                self.options.update(**kwargs)
+            else:
+                self.master.addons.invoke_addon_sync(addon, hooks.ConfigureHook(set()))
+
+    def script(self, path):
+        """
+        Loads a script from path, and returns the enclosed addon.
+        """
+        sc = script.Script(path, False)
+        return sc.addons[0] if sc.addons else None
+
+    def command(self, func, *args):
+        """
+        Invoke a command function with a list of string arguments within a command context, mimicking the actual command environment.
+        """
+        cmd = command.Command(self.master.commands, "test.command", func)
+        return cmd.call(args)

venv/Lib/site-packages/mitmproxy/test/tflow.py

@@ -0,0 +1,300 @@
+import time
+import uuid
+
+from wsproto.frame_protocol import Opcode
+
+from mitmproxy import connection
+from mitmproxy import dns
+from mitmproxy import flow
+from mitmproxy import http
+from mitmproxy import tcp
+from mitmproxy import udp
+from mitmproxy import websocket
+from mitmproxy.connection import ConnectionState
+from mitmproxy.proxy.mode_specs import ProxyMode
+from mitmproxy.test.tutils import tdnsreq
+from mitmproxy.test.tutils import tdnsresp
+from mitmproxy.test.tutils import treq
+from mitmproxy.test.tutils import tresp
+
+
+def ttcpflow(
+    client_conn=True, server_conn=True, messages=True, err=None
+) -> tcp.TCPFlow:
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if messages is True:
+        messages = [
+            tcp.TCPMessage(True, b"hello", 946681204.2),
+            tcp.TCPMessage(False, b"it's me", 946681204.5),
+        ]
+    if err is True:
+        err = terr()
+
+    f = tcp.TCPFlow(client_conn, server_conn)
+    f.timestamp_created = client_conn.timestamp_start
+    f.messages = messages
+    f.error = err
+    f.live = True
+    return f
+
+
+def tudpflow(
+    client_conn=True, server_conn=True, messages=True, err=None
+) -> udp.UDPFlow:
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if messages is True:
+        messages = [
+            udp.UDPMessage(True, b"hello", 946681204.2),
+            udp.UDPMessage(False, b"it's me", 946681204.5),
+        ]
+    if err is True:
+        err = terr()
+
+    f = udp.UDPFlow(client_conn, server_conn)
+    f.timestamp_created = client_conn.timestamp_start
+    f.messages = messages
+    f.error = err
+    f.live = True
+    return f
+
+
+def twebsocketflow(
+    messages=True, err=None, close_code=None, close_reason=""
+) -> http.HTTPFlow:
+    flow = http.HTTPFlow(tclient_conn(), tserver_conn())
+    flow.request = http.Request(
+        "example.com",
+        80,
+        b"GET",
+        b"http",
+        b"example.com",
+        b"/ws",
+        b"HTTP/1.1",
+        headers=http.Headers(
+            connection="upgrade",
+            upgrade="websocket",
+            sec_websocket_version="13",
+            sec_websocket_key="1234",
+        ),
+        content=b"",
+        trailers=None,
+        timestamp_start=946681200,
+        timestamp_end=946681201,
+    )
+    flow.response = http.Response(
+        b"HTTP/1.1",
+        101,
+        reason=b"Switching Protocols",
+        headers=http.Headers(
+            connection="upgrade",
+            upgrade="websocket",
+            sec_websocket_accept=b"",
+        ),
+        content=b"",
+        trailers=None,
+        timestamp_start=946681202,
+        timestamp_end=946681203,
+    )
+
+    flow.websocket = twebsocket()
+
+    flow.websocket.close_reason = close_reason
+
+    if close_code is not None:
+        flow.websocket.close_code = close_code
+    else:
+        if err is True:
+            # ABNORMAL_CLOSURE
+            flow.websocket.close_code = 1006
+        else:
+            # NORMAL_CLOSURE
+            flow.websocket.close_code = 1000
+
+    flow.live = True
+    return flow
+
+
+def tdnsflow(
+    *,
+    client_conn: connection.Client | None = None,
+    server_conn: connection.Server | None = None,
+    req: dns.DNSMessage | None = None,
+    resp: bool | dns.DNSMessage = False,
+    err: bool | flow.Error = False,
+    live: bool = True,
+) -> dns.DNSFlow:
+    """Create a DNS flow for testing."""
+    if client_conn is None:
+        client_conn = tclient_conn()
+        client_conn.proxy_mode = ProxyMode.parse("dns")
+        client_conn.transport_protocol = "udp"
+    if server_conn is None:
+        server_conn = tserver_conn()
+        server_conn.transport_protocol = "udp"
+    if req is None:
+        req = tdnsreq()
+
+    if resp is True:
+        resp = tdnsresp()
+    if err is True:
+        err = terr()
+
+    assert resp is False or isinstance(resp, dns.DNSMessage)
+    assert err is False or isinstance(err, flow.Error)
+
+    f = dns.DNSFlow(client_conn, server_conn)
+    f.timestamp_created = req.timestamp or time.time()
+    f.request = req
+    f.response = resp or None
+    f.error = err or None
+    f.live = live
+    return f
+
+
+def tflow(
+    *,
+    client_conn: connection.Client | None = None,
+    server_conn: connection.Server | None = None,
+    req: http.Request | None = None,
+    resp: bool | http.Response = False,
+    err: bool | flow.Error = False,
+    ws: bool | websocket.WebSocketData = False,
+    live: bool = True,
+) -> http.HTTPFlow:
+    """Create a flow for testing."""
+    if client_conn is None:
+        client_conn = tclient_conn()
+    if server_conn is None:
+        server_conn = tserver_conn()
+    if req is None:
+        req = treq()
+
+    if resp is True:
+        resp = tresp()
+    if err is True:
+        err = terr()
+    if ws is True:
+        ws = twebsocket()
+
+    assert resp is False or isinstance(resp, http.Response)
+    assert err is False or isinstance(err, flow.Error)
+    assert ws is False or isinstance(ws, websocket.WebSocketData)
+
+    f = http.HTTPFlow(client_conn, server_conn)
+    f.timestamp_created = req.timestamp_start
+    f.request = req
+    f.response = resp or None
+    f.error = err or None
+    f.websocket = ws or None
+    f.live = live
+    return f
+
+
+class DummyFlow(flow.Flow):
+    """A flow that is neither HTTP nor TCP."""
+
+
+def tdummyflow(client_conn=True, server_conn=True, err=None) -> DummyFlow:
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if err is True:
+        err = terr()
+
+    f = DummyFlow(client_conn, server_conn)
+    f.error = err
+    f.live = True
+    return f
+
+
+def tclient_conn() -> connection.Client:
+    c = connection.Client(
+        id=str(uuid.uuid4()),
+        peername=("127.0.0.1", 22),
+        sockname=("", 0),
+        mitmcert=None,
+        timestamp_start=946681200,
+        timestamp_tls_setup=946681201,
+        timestamp_end=946681206,
+        sni="address",
+        cipher="cipher",
+        alpn=b"http/1.1",
+        tls_version="TLSv1.2",
+        state=ConnectionState.OPEN,
+        error=None,
+        tls=False,
+        certificate_list=[],
+        alpn_offers=[],
+        cipher_list=[],
+        proxy_mode=ProxyMode.parse("regular"),
+    )
+    return c
+
+
+def tserver_conn() -> connection.Server:
+    c = connection.Server(
+        id=str(uuid.uuid4()),
+        address=("address", 22),
+        peername=("192.168.0.1", 22),
+        sockname=("address", 22),
+        timestamp_start=946681202,
+        timestamp_tcp_setup=946681203,
+        timestamp_tls_setup=946681204,
+        timestamp_end=946681205,
+        sni="address",
+        alpn=None,
+        tls_version="TLSv1.2",
+        via=None,
+        state=ConnectionState.CLOSED,
+        error=None,
+        tls=False,
+        certificate_list=[],
+        alpn_offers=[],
+        cipher=None,
+        cipher_list=[],
+    )
+    return c
+
+
+def terr(content: str = "error") -> flow.Error:
+    err = flow.Error(content, 946681207)
+    return err
+
+
+def twebsocket(messages: bool = True) -> websocket.WebSocketData:
+    ws = websocket.WebSocketData()
+
+    if messages:
+        ws.messages = [
+            websocket.WebSocketMessage(Opcode.BINARY, True, b"hello binary", 946681203),
+            websocket.WebSocketMessage(Opcode.TEXT, True, b"hello text", 946681204),
+            websocket.WebSocketMessage(Opcode.TEXT, False, b"it's me", 946681205),
+        ]
+    ws.close_reason = "Close Reason"
+    ws.close_code = 1000
+    ws.closed_by_client = False
+    ws.timestamp_end = 946681205
+
+    return ws
+
+
+def tflows() -> list[flow.Flow]:
+    return [
+        tflow(resp=True),
+        tflow(err=True),
+        tflow(ws=True),
+        ttcpflow(),
+        ttcpflow(err=True),
+        tudpflow(),
+        tudpflow(err=True),
+        tdnsflow(resp=True),
+        tdnsflow(req=tdnsreq(questions=[])),
+        tdnsflow(err=True),
+    ]

venv/Lib/site-packages/mitmproxy/test/tutils.py

@@ -0,0 +1,95 @@
+from mitmproxy import dns
+from mitmproxy import http
+
+
+def tdnsreq(**kwargs) -> dns.DNSMessage:
+    default = dict(
+        timestamp=946681200,
+        id=42,
+        query=True,
+        op_code=dns.op_codes.QUERY,
+        authoritative_answer=False,
+        truncation=False,
+        recursion_desired=True,
+        recursion_available=False,
+        reserved=0,
+        response_code=dns.response_codes.NOERROR,
+        questions=[dns.Question("dns.google", dns.types.A, dns.classes.IN)],
+        answers=[],
+        authorities=[],
+        additionals=[],
+    )
+    default.update(kwargs)
+    return dns.DNSMessage(**default)  # type: ignore
+
+
+def tdnsresp(**kwargs) -> dns.DNSMessage:
+    default = dict(
+        timestamp=946681201,
+        id=42,
+        query=False,
+        op_code=dns.op_codes.QUERY,
+        authoritative_answer=False,
+        truncation=False,
+        recursion_desired=True,
+        recursion_available=True,
+        reserved=0,
+        response_code=dns.response_codes.NOERROR,
+        questions=[dns.Question("dns.google", dns.types.A, dns.classes.IN)],
+        answers=[
+            dns.ResourceRecord(
+                "dns.google", dns.types.A, dns.classes.IN, 32, b"\x08\x08\x08\x08"
+            ),
+            dns.ResourceRecord(
+                "dns.google", dns.types.A, dns.classes.IN, 32, b"\x08\x08\x04\x04"
+            ),
+        ],
+        authorities=[],
+        additionals=[],
+    )
+    default.update(kwargs)
+    return dns.DNSMessage(**default)  # type: ignore
+
+
+def treq(**kwargs) -> http.Request:
+    """
+    Returns:
+        mitmproxy.net.http.Request
+    """
+    default = dict(
+        host="address",
+        port=22,
+        method=b"GET",
+        scheme=b"http",
+        authority=b"",
+        path=b"/path",
+        http_version=b"HTTP/1.1",
+        headers=http.Headers(((b"header", b"qvalue"), (b"content-length", b"7"))),
+        content=b"content",
+        trailers=None,
+        timestamp_start=946681200,
+        timestamp_end=946681201,
+    )
+    default.update(kwargs)
+    return http.Request(**default)  # type: ignore
+
+
+def tresp(**kwargs) -> http.Response:
+    """
+    Returns:
+        mitmproxy.net.http.Response
+    """
+    default = dict(
+        http_version=b"HTTP/1.1",
+        status_code=200,
+        reason=b"OK",
+        headers=http.Headers(
+            ((b"header-response", b"svalue"), (b"content-length", b"7"))
+        ),
+        content=b"message",
+        trailers=None,
+        timestamp_start=946681202,
+        timestamp_end=946681203,
+    )
+    default.update(kwargs)
+    return http.Response(**default)  # type: ignore