dataclasses, further development

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 = ""
+    tls: str
-    content = ""
+    content: str
-    scheme = ""
+    scheme: str
-    method = ""
+    method: str
-    host = ""
+    host: str
-    port = 0
+    port: int
-    http_version = ""
+    http_version: str
-    path = ""
+    path: str
-    timestamp_start = 0.0
+    timestamp_start: float
-    timestamp_end = 0.0
+    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
+    status_code: int
-    http_version = ""
+    http_version: str
-    reason = ""
+    reason: str
-    content = ""
+    content: str
-    timestamp_start = 0.0
+    timestamp_start: float
-    timestamp_end = 0.0
+    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 = "";
+    uid: str;
-    request = None
+    request: bRequest
-    response = None
+    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
+        # send a single message, no checks involved
-        def update_flows(self):
+        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):
+                    if self.flows.get(i, None):
-                        print("flow {} doubled? ignoring...")
+                        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
+        # update all current flows
-        def handle_flow(self, state, flow):
+        def update_flows(self):
-            for i, v in flows.items:
+            for k,v in self.flows.items():
                 state, flow, timer, retries = v
 
+                # state machine for flows
+
                 if state == FlowState.UNSENT_REQ:
-                    # send it
+                    msg = b""
-                    a = convert_to_strings(msg)
+                    # send the request
-                    self.socket.send(str.encode(json.dumps(a)))
+                    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"})
+                        pass
-                        
+                        #self.send_msg_and_ack({"msg": "ping"})
-        def send(self, msg):
-            a = convert_to_strings(msg)
-            self.socket.send(str.encode(json.dumps(a)))
 
         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")
+                        #print("m sending")
-                        a = convert_to_strings(msg)
+                        self.send(msg)
-                        self.socket.send(str.encode(json.dumps(a)))
                         if (self.socket.poll(50) & zmq.POLLIN) != 0:
                                 msg = self.socket.recv()
                                 try: