bigsnitch

mitmaddon/bigsnitch.py

@@ -1,6 +1,9 @@
 #!/usr/bin/env python3
+#
+
 import pdb
-from mitmproxy import ctx
+
+from mitmproxy import ctx, Flow
 import threading
 from queue import Queue, Empty
 import time
@@ -8,7 +11,7 @@ import zmq
 import json
 from enum import Enum
 from dataclasses import dataclass
-from typing import List
+from typing import List, Dict
 
 # this method is used to convert flow states (generated with get_state()) to json
 def convert_to_strings(obj):
@@ -55,6 +58,7 @@ class bRequest:
 
     error: str
 
+    # init from dict
     def __init__(self, flow: dict):
         self.server_ip_address = flow["server_ip_address"]
         self.tls = flow["server_conn"]["tls_established"]
@@ -103,28 +107,53 @@ class bFlow:
         self.request =  bRequest(flow["request"])
         self.response =  bRequest(flow["response"])
 
+#
+# Networkthread state machine types
+#
+
+@dataclass
 class FlowState(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
+    REQUEST = 2
+    RESPONSEHEADERS = 3
+    RESPONSE = 4
+
+# for use in NetworkThread queue
+@dataclass
+class FlowItem:
+    id: int
+    mitmstate: MitmState
+    state: FlowState
+    flow: Flow
+    time: float = 0
+    retries_left: int = 5
+
 """
 
 The network thread communicates with the bigsnitch plugin using zeromq.
 
 """
 
+@dataclass
 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
+                # for zmq use
-                # id : (state, flow, timer, retries left)
-                self.flows = {}
                 self.context = zmq.Context()
+                # all current flows being handled by mitmproxy
+                self.flows: Dict[FlowItem]
                 # timer for sending pings to check if the connection broke
                 self.timer = time.monotonic()
                 # retries left for reconnecting / resending a broken flow
@@ -147,9 +176,9 @@ class NetworkThread(threading.Thread):
 
                     # csave the new flows, if necessary
                     if typ == "request":
-                        self.flows[i] = (FlowState.UNSENT_REQ, flow, time.monotonic(), 5)
+                        self.flows[i] = bFlow(FlowState.UNSENT_REQ, flow, time.monotonic(), 5)
                     elif typ == "response":
-                        self.flows[i] = (FlowState.UNSENT_RES, flow, time.monotonic(), 5)
+                        self.flows[i] = bFlow(FlowState.UNSENT_RES, flow, time.monotonic(), 5)
                 except Empty:
                     break
 
@@ -164,7 +193,6 @@ class NetworkThread(threading.Thread):
                     msg = b""
                     # send the request
                     self.send(msg)
-                    pass
                 elif state == FlowState.SENT_REQ:
                     # check timer, try resend
                     pass