changes

mitmaddon/bigsnitch.py

@@ -22,7 +22,6 @@ def convert_to_strings(obj):
         elif isinstance(obj, list) or isinstance(obj, tuple):
             return [convert_to_strings(element) for element in obj]
         elif isinstance(obj, bytes):
-            data = ""
             try:
                 data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
             except:
@@ -99,41 +98,45 @@ class bResponse:
 
 @dataclass
 class bFlow:
-    uid: str;
+    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"])
+        self.response =  bResponse(flow["response"])
 
-#
-# Networkthread state machine types
-#
 
 @dataclass
-class FlowState(Enum):
+class bFlowState(Enum):
-    UNSENT_REQ = 0
-    SENT_REQ = 1
-    UNSENT_RES = 2
-    SENT_RES = 3
-
-# current flow state in Mitmproxy
-@dataclass
-class MitmState(Enum):
     ERROR = 0
-    REQUESTHEADERS = 1
+    UNSENT_HTTP_REQUEST = 1
-    REQUEST = 2
+    SENT_HTTP_REQUEST = 2
-    RESPONSEHEADERS = 3
+    UNSENT_HTTP_RESPONSE = 3
-    RESPONSE = 4
+    SENT_HTTP_RESPONSE = 4
+
+@dataclass
+class bPacketType:
+    NACK = 0
+    ACK = 1
+    KILL = 2
+    WARNING = 3
+    ERROR = 4
+    PING = 5
+    HTTP_REQUEST = 6
+    HTTP_RESPONSE = 7
+
+@dataclass
+class bPacket:
+    ptype: bPacketType
+    flowid: int
+    data: str
 
 # for use in NetworkThread queue
 @dataclass
 class FlowItem:
-    id: int
+    state: bFlowState
-    mitmstate: MitmState
-    state: FlowState
     flow: Flow
     time: float = 0
     retries_left: int = 5
@@ -145,10 +148,11 @@ The network thread communicates with the bigsnitch plugin using zeromq.
 """
 
 class NetworkThread(threading.Thread):
-        def __init__(self, name, queue):
+        def __init__(self, name: str, queue: Queue):
                 threading.Thread.__init__(self)
                 self.name = name
                 # queue for communicating with the main mitmproxy thread
+                # contains tuples of (id, FlowItem)
                 self.q = queue
                 # for zmq use
                 self.context = zmq.Context()
@@ -169,38 +173,72 @@ class NetworkThread(threading.Thread):
             while True:
                 try:
                     # get new flows that may occured
-                    i, flow, typ = self.q.get(block=False)
+                    i, flowitem = self.q.get(block=False)
                     if self.flows.get(i, None):
                         print(f"flow {i} doubled? ignoring...")
                         continue
+                    else:
+                        self.flows.append(flowitem)
 
-                    # csave the new flows, if necessary
-                    if typ == "request":
-                        self.flows[i] = bFlow(FlowState.UNSENT_REQ, flow, time.monotonic(), 5)
-                    elif typ == "response":
-                        self.flows[i] = bFlow(FlowState.UNSENT_RES, flow, time.monotonic(), 5)
                 except Empty:
                     break
 
+        def send_packet(self, pkg: bPacket):
+            msg = {"type": pkg.ptype, "id": pkg.flowid, "data": pkg.data}
+            self.send(msg)
+
+        def send_http_request(self, id: int, request: bRequest):
+            pkg = bPacket(bPacketType.HTTP_REQUEST, id, request.json())
+            self.send_packet(pkg)
+            flows[id].state = bFlowState.SENT_HTTP_REQUEST
+            flows[id].time = time.monotonic()
+
+        def send_http_response(self, id: int, response: bResponse):
+            pkg = bPacket(bPacketType.HTTP_RESPONSE, id, response.json())
+            self.send_packet(pkg)
+            flows[id].state = bFlowState.SENT_HTTP_RESPONSE
+            flows[id].time = time.monotonic()
+
         # update all current flows
+        # handles the state machine for each flow
         def update_flows(self):
-            for k,v in self.flows.items():
+            for id, flow in self.flows.items():
-                state, flow, timer, retries = v
+                # send the request if not sent
+                if state == bFlowState.UNSENT_HTTP_REQUEST:
+                    self.send_request(id, bRequest(flow.flow))
 
-                # state machine for flows
+                elif state == bFlowState.SENT_HTTP_REQUEST:
-
-                if state == FlowState.UNSENT_REQ:
-                    msg = b""
-                    # send the request
-                    self.send(msg)
-                elif state == FlowState.SENT_REQ:
                     # check timer, try resend
-                    pass
+                    delta = time.monotonic() - flows[id].time
-                elif state == FlowState.UNSENT_RES:
+                    if delta > 5:
-                    pass
+                        self.send_request(id, bRequest(flow.flow))
-                elif state == FlowState.SENT_RES:
+                        flows[id].retries -= 1
-                    pass
+                        continue
 
+                    if flows[id].retries <= 0:
+                        flows[id].flow.kill()
+                        print(f"http request {id} timed out! flow killed.")
+
+                # send the response if not sent
+                elif state == bFlowState.UNSENT_HTTP_RESPONSE:
+                    self.send_response(id, bResponse(flow.flow))
+
+                elif state == bFlowState.SENT_HTTP_RESPONSE:
+                    # check timer, try resend
+                    delta = time.monotonic() - flows[id].time
+                    if delta > 5:
+                        self.send_response(id, bResponse(flow.flow))
+                        flows[id].retries -= 1
+                        continue
+
+                    if flows[id].retries <= 0:
+                        flows[id].flow.kill()
+                        print(f"http response {id} timedout! flow killed.")
+
+                elif state == bFlowState.ERROR:
+                    print(f"error in flow {id}!")
+
+        # handle incoming packets / update the statemachine
         def handle_packets(self):
             while((self.socket.poll(50) & zmq.POLLIN) != 0):
                 msg = self.socket.recv()
@@ -242,9 +280,10 @@ class NetworkThread(threading.Thread):
         def connect(self):
                 self.socket = self.context.socket(zmq.PAIR)
                 self.socket.connect("tcp://127.0.0.1:12345")
-                self.send_msg_and_ack({"msg": "ping"})
+                #self.send_msg_and_ack({"msg": "ping"})
-                print("successfully connected")
+                print("connected")
 
+        """
         def send_msg_and_ack(self, msg):
                 self.timer = time.monotonic()
                 while True:
@@ -266,6 +305,7 @@ class NetworkThread(threading.Thread):
                         print("no ack received, reconnecting...")
                         self.reconnect()
                 return NO_MSG
+        """
 
 class BigSnitchBridge:
     def __init__(self):
@@ -274,17 +314,20 @@ class BigSnitchBridge:
         self.thread.start()
     
     def request(self, flow):
-        self.q.put_nowait((flow.id, flow, "request"))
+        flowitem = FlowItem(bFlowState.UNSENT_HTTP_REQUEST, flow)
+        self.q.put_nowait((flow.id, flowitem))
         # intercept until ACK received
         flow.intercept()
 
     def response(self, flow):
-        self.q.put_nowait((flow.id, flow, "response"))
+        flowitem = FlowItem(bFlowState.UNSENT_HTTP_RESPONSE, flow)
+        self.q.put_nowait((flow.id, flowitem))
         # intercept until ACK received
         flow.intercept()
 
     def error(self, flow):
-        self.q.put_nowait((flow.id, flow, "error"))
+        flowitem = FlowItem(bFlowState.ERROR, flow, time.monotonic())
+        self.q.put_nowait((flow.id, flowitem))
 
 
 """