From 8a784f9c481b1b047950a523f27989d95ba3ca90 Mon Sep 17 00:00:00 2001 From: end Date: Thu, 22 Apr 2021 00:41:13 +0200 Subject: [PATCH] dataclasses, further development --- mitmaddon/bigsnitch.py | 150 +++++++++++++++++++++++++++-------------- 1 file changed, 100 insertions(+), 50 deletions(-) diff --git a/mitmaddon/bigsnitch.py b/mitmaddon/bigsnitch.py index b423250..6ee3e9a 100644 --- a/mitmaddon/bigsnitch.py +++ b/mitmaddon/bigsnitch.py @@ -7,6 +7,8 @@ import time import zmq import json from enum import Enum +from dataclasses import dataclass +from typing import List # this method is used to convert flow states (generated with get_state()) to json def convert_to_strings(obj): @@ -25,46 +27,87 @@ def convert_to_strings(obj): return data return obj +# +# bigsnitch communication types +# -# bigsnitch Request type +@dataclass +class bHeader: + key: str + value: str + +@dataclass class bRequest: - server_ip_address = "" + server_ip_address: str - tls = "" - content = "" - scheme = "" - method = "" - host = "" - port = 0 - http_version = "" - path = "" - timestamp_start = 0.0 - timestamp_end = 0.0 + tls: str + content: str + scheme: str + method: str + host: str + port: int + http_version: str + path: str + timestamp_start: float + timestamp_end: float # [("Header","Data")] - headers = [] + headers: List[bHeader] - error = "" + error: str + def __init__(self, flow: dict): + self.server_ip_address = flow["server_ip_address"] + self.tls = flow["server_conn"]["tls_established"] + self.content = flow["content"] + self.scheme = flow["scheme"] + self.method = flow["method"] + self.host = flow["host"] + self.port = flow["port"] + self.http_version = flow["http_version"] + self.timestamp_start = flow["timestamp_start"] + self.timestamp_end = flow["timestamp_end"] + + for k,v in flow["headers"]: + self.headers.append(bHeader(k,v)) + +@dataclass class bResponse: - status_code = 0 - http_version = "" - reason = "" - content = "" - timestamp_start = 0.0 - timestamp_end = 0.0 + status_code: int + http_version: str + reason: str + content: str + timestamp_start: float + timestamp_end: float # [("Header","Data")] - headers = [] + headers: List[bHeader] + def __init__(self, flow: dict): + self.status_code = flow["status_code"] + self.http_version = flow["http_version"] + self.reason = flow["reason"] + self.content = flow["content"] + self.timestamp_start = flow["timestamp_start"] + self.timestamp_end = flow["timestamp_end"] + + for k,v in flow["headers"]: + self.headers.append(bHeader(k,v)) + +@dataclass class bFlow: - uid = ""; - request = None - response = None + uid: str; + request: bRequest + response: bResponse + + def __init__(self, flow: dict): + self.uid = flow["id"] + self.request = bRequest(flow["request"]) + self.response = bRequest(flow["response"]) class FlowState(Enum): UNSENT_REQ = 0 SENT_REQ = 1 UNSENT_RES = 2 - SENT_REQ = 3 + SENT_RES = 3 """ @@ -76,40 +119,51 @@ class NetworkThread(threading.Thread): def __init__(self, name, queue): threading.Thread.__init__(self) self.name = name + # queue for communicating with the main mitmproxy thread self.q = queue + # all current flows being handled by mitmproxy # id : (state, flow, timer, retries left) self.flows = {} self.context = zmq.Context() - # timer for sending pings + # timer for sending pings to check if the connection broke self.timer = time.monotonic() - # retries left for reconnecting + # retries left for reconnecting / resending a broken flow self.retries = 5 - # get new flows that may occured - def update_flows(self): + # send a single message, no checks involved + def send(self, msg): + a = convert_to_strings(msg) + self.socket.send(str.encode(json.dumps(a))) + + # add new flows from the queue + def get_new_flows(self): while True: try: + # get new flows that may occured i, flow, typ = self.q.get(block=False) - if flows.get(i, None): - print("flow {} doubled? ignoring...") + if self.flows.get(i, None): + print(f"flow {i} doubled? ignoring...") continue + # csave the new flows, if necessary if typ == "request": - flows[i] = (FlowState.UNSENT_REQ, flow, self.timer, 5) + self.flows[i] = (FlowState.UNSENT_REQ, flow, time.monotonic(), 5) elif typ == "response": - flows[i] = (FlowState.UNSENT_RES, flow, self.timer, 5) + self.flows[i] = (FlowState.UNSENT_RES, flow, time.monotonic(), 5) except Empty: break - # state machine for flows - def handle_flow(self, state, flow): - for i, v in flows.items: + # update all current flows + def update_flows(self): + for k,v in self.flows.items(): state, flow, timer, retries = v + # state machine for flows + if state == FlowState.UNSENT_REQ: - # send it - a = convert_to_strings(msg) - self.socket.send(str.encode(json.dumps(a))) + msg = b"" + # send the request + self.send(msg) pass elif state == FlowState.SENT_REQ: # check timer, try resend @@ -119,8 +173,6 @@ class NetworkThread(threading.Thread): elif state == FlowState.SENT_RES: pass - self.send_msg_and_ack(a) - def handle_packets(self): while((self.socket.poll(50) & zmq.POLLIN) != 0): msg = self.socket.recv() @@ -141,14 +193,13 @@ class NetworkThread(threading.Thread): self.connect() while True: self.timer = time.monotonic() - update_flows() + self.get_new_flows() + self.handle_packets() + self.update_flows() if self.timer - time.monotonic() < -5: - self.send_msg_and_ack({"msg": "ping"}) - - def send(self, msg): - a = convert_to_strings(msg) - self.socket.send(str.encode(json.dumps(a))) + pass + #self.send_msg_and_ack({"msg": "ping"}) def disconnect(self): self.socket.setsockopt(zmq.LINGER,0) @@ -169,9 +220,8 @@ class NetworkThread(threading.Thread): def send_msg_and_ack(self, msg): self.timer = time.monotonic() while True: - print("m sending") - a = convert_to_strings(msg) - self.socket.send(str.encode(json.dumps(a))) + #print("m sending") + self.send(msg) if (self.socket.poll(50) & zmq.POLLIN) != 0: msg = self.socket.recv() try: