codekingpro/portable-devtools

1import time2import uuid3 4from wsproto.frame_protocol import Opcode5 6from mitmproxy import connection7from mitmproxy import dns8from mitmproxy import flow9from mitmproxy import http10from mitmproxy import tcp11from mitmproxy import udp12from mitmproxy import websocket13from mitmproxy.connection import ConnectionState14from mitmproxy.proxy.mode_specs import ProxyMode15from mitmproxy.test.tutils import tdnsreq16from mitmproxy.test.tutils import tdnsresp17from mitmproxy.test.tutils import treq18from mitmproxy.test.tutils import tresp19 20 21def ttcpflow(22    client_conn=True, server_conn=True, messages=True, err=None23) -> tcp.TCPFlow:24    if client_conn is True:25        client_conn = tclient_conn()26    if server_conn is True:27        server_conn = tserver_conn()28    if messages is True:29        messages = [30            tcp.TCPMessage(True, b"hello", 946681204.2),31            tcp.TCPMessage(False, b"it's me", 946681204.5),32        ]33    if err is True:34        err = terr()35 36    f = tcp.TCPFlow(client_conn, server_conn)37    f.timestamp_created = client_conn.timestamp_start38    f.messages = messages39    f.error = err40    f.live = True41    return f42 43 44def tudpflow(45    client_conn=True, server_conn=True, messages=True, err=None46) -> udp.UDPFlow:47    if client_conn is True:48        client_conn = tclient_conn()49    if server_conn is True:50        server_conn = tserver_conn()51    if messages is True:52        messages = [53            udp.UDPMessage(True, b"hello", 946681204.2),54            udp.UDPMessage(False, b"it's me", 946681204.5),55        ]56    if err is True:57        err = terr()58 59    f = udp.UDPFlow(client_conn, server_conn)60    f.timestamp_created = client_conn.timestamp_start61    f.messages = messages62    f.error = err63    f.live = True64    return f65 66 67def twebsocketflow(68    messages=True, err=None, close_code=None, close_reason=""69) -> http.HTTPFlow:70    flow = http.HTTPFlow(tclient_conn(), tserver_conn())71    flow.request = http.Request(72        "example.com",73        80,74        b"GET",75        b"http",76        b"example.com",77        b"/ws",78        b"HTTP/1.1",79        headers=http.Headers(80            connection="upgrade",81            upgrade="websocket",82            sec_websocket_version="13",83            sec_websocket_key="1234",84        ),85        content=b"",86        trailers=None,87        timestamp_start=946681200,88        timestamp_end=946681201,89    )90    flow.response = http.Response(91        b"HTTP/1.1",92        101,93        reason=b"Switching Protocols",94        headers=http.Headers(95            connection="upgrade",96            upgrade="websocket",97            sec_websocket_accept=b"",98        ),99        content=b"",100        trailers=None,101        timestamp_start=946681202,102        timestamp_end=946681203,103    )104 105    flow.websocket = twebsocket()106 107    flow.websocket.close_reason = close_reason108 109    if close_code is not None:110        flow.websocket.close_code = close_code111    else:112        if err is True:113            # ABNORMAL_CLOSURE114            flow.websocket.close_code = 1006115        else:116            # NORMAL_CLOSURE117            flow.websocket.close_code = 1000118 119    flow.live = True120    return flow121 122 123def tdnsflow(124    *,125    client_conn: connection.Client | None = None,126    server_conn: connection.Server | None = None,127    req: dns.DNSMessage | None = None,128    resp: bool | dns.DNSMessage = False,129    err: bool | flow.Error = False,130    live: bool = True,131) -> dns.DNSFlow:132    """Create a DNS flow for testing."""133    if client_conn is None:134        client_conn = tclient_conn()135        client_conn.proxy_mode = ProxyMode.parse("dns")136        client_conn.transport_protocol = "udp"137    if server_conn is None:138        server_conn = tserver_conn()139        server_conn.transport_protocol = "udp"140    if req is None:141        req = tdnsreq()142 143    if resp is True:144        resp = tdnsresp()145    if err is True:146        err = terr()147 148    assert resp is False or isinstance(resp, dns.DNSMessage)149    assert err is False or isinstance(err, flow.Error)150 151    f = dns.DNSFlow(client_conn, server_conn)152    f.timestamp_created = req.timestamp or time.time()153    f.request = req154    f.response = resp or None155    f.error = err or None156    f.live = live157    return f158 159 160def tflow(161    *,162    client_conn: connection.Client | None = None,163    server_conn: connection.Server | None = None,164    req: http.Request | None = None,165    resp: bool | http.Response = False,166    err: bool | flow.Error = False,167    ws: bool | websocket.WebSocketData = False,168    live: bool = True,169) -> http.HTTPFlow:170    """Create a flow for testing."""171    if client_conn is None:172        client_conn = tclient_conn()173    if server_conn is None:174        server_conn = tserver_conn()175    if req is None:176        req = treq()177 178    if resp is True:179        resp = tresp()180    if err is True:181        err = terr()182    if ws is True:183        ws = twebsocket()184 185    assert resp is False or isinstance(resp, http.Response)186    assert err is False or isinstance(err, flow.Error)187    assert ws is False or isinstance(ws, websocket.WebSocketData)188 189    f = http.HTTPFlow(client_conn, server_conn)190    f.timestamp_created = req.timestamp_start191    f.request = req192    f.response = resp or None193    f.error = err or None194    f.websocket = ws or None195    f.live = live196    return f197 198 199class DummyFlow(flow.Flow):200    """A flow that is neither HTTP nor TCP."""201 202 203def tdummyflow(client_conn=True, server_conn=True, err=None) -> DummyFlow:204    if client_conn is True:205        client_conn = tclient_conn()206    if server_conn is True:207        server_conn = tserver_conn()208    if err is True:209        err = terr()210 211    f = DummyFlow(client_conn, server_conn)212    f.error = err213    f.live = True214    return f215 216 217def tclient_conn() -> connection.Client:218    c = connection.Client(219        id=str(uuid.uuid4()),220        peername=("127.0.0.1", 22),221        sockname=("", 0),222        mitmcert=None,223        timestamp_start=946681200,224        timestamp_tls_setup=946681201,225        timestamp_end=946681206,226        sni="address",227        cipher="cipher",228        alpn=b"http/1.1",229        tls_version="TLSv1.2",230        state=ConnectionState.OPEN,231        error=None,232        tls=False,233        certificate_list=[],234        alpn_offers=[],235        cipher_list=[],236        proxy_mode=ProxyMode.parse("regular"),237    )238    return c239 240 241def tserver_conn() -> connection.Server:242    c = connection.Server(243        id=str(uuid.uuid4()),244        address=("address", 22),245        peername=("192.168.0.1", 22),246        sockname=("address", 22),247        timestamp_start=946681202,248        timestamp_tcp_setup=946681203,249        timestamp_tls_setup=946681204,250        timestamp_end=946681205,251        sni="address",252        alpn=None,253        tls_version="TLSv1.2",254        via=None,255        state=ConnectionState.CLOSED,256        error=None,257        tls=False,258        certificate_list=[],259        alpn_offers=[],260        cipher=None,261        cipher_list=[],262    )263    return c264 265 266def terr(content: str = "error") -> flow.Error:267    err = flow.Error(content, 946681207)268    return err269 270 271def twebsocket(messages: bool = True) -> websocket.WebSocketData:272    ws = websocket.WebSocketData()273 274    if messages:275        ws.messages = [276            websocket.WebSocketMessage(Opcode.BINARY, True, b"hello binary", 946681203),277            websocket.WebSocketMessage(Opcode.TEXT, True, b"hello text", 946681204),278            websocket.WebSocketMessage(Opcode.TEXT, False, b"it's me", 946681205),279        ]280    ws.close_reason = "Close Reason"281    ws.close_code = 1000282    ws.closed_by_client = False283    ws.timestamp_end = 946681205284 285    return ws286 287 288def tflows() -> list[flow.Flow]:289    return [290        tflow(resp=True),291        tflow(err=True),292        tflow(ws=True),293        ttcpflow(),294        ttcpflow(err=True),295        tudpflow(),296        tudpflow(err=True),297        tdnsflow(resp=True),298        tdnsflow(req=tdnsreq(questions=[])),299        tdnsflow(err=True),300    ]301