codekingpro/portable-devtools

1from collections.abc import Callable2from collections.abc import Iterator3from typing import Any4 5from mitmproxy import dns6from mitmproxy import flow7from mitmproxy import hooks8from mitmproxy import http9from mitmproxy import tcp10from mitmproxy import udp11from mitmproxy.proxy import layers12 13TEventGenerator = Iterator[hooks.Hook]14 15 16def _iterate_http(f: http.HTTPFlow) -> TEventGenerator:17    if f.request:18        yield layers.http.HttpRequestHeadersHook(f)19        yield layers.http.HttpRequestHook(f)20    if f.response:21        yield layers.http.HttpResponseHeadersHook(f)22        yield layers.http.HttpResponseHook(f)23    if f.websocket:24        message_queue = f.websocket.messages25        f.websocket.messages = []26        yield layers.websocket.WebsocketStartHook(f)27        for m in message_queue:28            f.websocket.messages.append(m)29            yield layers.websocket.WebsocketMessageHook(f)30        yield layers.websocket.WebsocketEndHook(f)31    elif f.error:32        yield layers.http.HttpErrorHook(f)33 34 35def _iterate_tcp(f: tcp.TCPFlow) -> TEventGenerator:36    messages = f.messages37    f.messages = []38    yield layers.tcp.TcpStartHook(f)39    while messages:40        f.messages.append(messages.pop(0))41        yield layers.tcp.TcpMessageHook(f)42    if f.error:43        yield layers.tcp.TcpErrorHook(f)44    else:45        yield layers.tcp.TcpEndHook(f)46 47 48def _iterate_udp(f: udp.UDPFlow) -> TEventGenerator:49    messages = f.messages50    f.messages = []51    yield layers.udp.UdpStartHook(f)52    while messages:53        f.messages.append(messages.pop(0))54        yield layers.udp.UdpMessageHook(f)55    if f.error:56        yield layers.udp.UdpErrorHook(f)57    else:58        yield layers.udp.UdpEndHook(f)59 60 61def _iterate_dns(f: dns.DNSFlow) -> TEventGenerator:62    if f.request:63        yield layers.dns.DnsRequestHook(f)64    if f.response:65        yield layers.dns.DnsResponseHook(f)66    if f.error:67        yield layers.dns.DnsErrorHook(f)68 69 70_iterate_map: dict[type[flow.Flow], Callable[[Any], TEventGenerator]] = {71    http.HTTPFlow: _iterate_http,72    tcp.TCPFlow: _iterate_tcp,73    udp.UDPFlow: _iterate_udp,74    dns.DNSFlow: _iterate_dns,75}76 77 78def iterate(f: flow.Flow) -> TEventGenerator:79    try:80        e = _iterate_map[type(f)]81    except KeyError as err:82        raise TypeError(f"Unknown flow type: {f.__class__.__name__}") from err83    else:84        yield from e(f)85