foo

mitmaddon/bigsnitch.py

@@ -3,312 +3,12 @@
 
 import pdb
 
-from mitmproxy import ctx
-from mitmproxy.flow import Flow
-import threading
 from queue import Queue, Empty
-import time
+from networkthread import bFlowState, FlowItem, NetworkThread
-import zmq
-import json
-from enum import Enum
-from dataclasses import dataclass
-from typing import List, Dict
-
-# this method is used to convert flow states (generated with get_state()) to json
-def convert_to_strings(obj):
-        if isinstance(obj, dict):
-            return {convert_to_strings(key): convert_to_strings(value)
-                    for key, value in obj.items()}
-        elif isinstance(obj, list) or isinstance(obj, tuple):
-            return [convert_to_strings(element) for element in obj]
-        elif isinstance(obj, bytes):
-            try:
-                data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
-            except:
-                print(obj)
-                data = str(obj)[2:-1]
-            return data
-
-        return obj
-#
-# bigsnitch communication types
-#
-
-@dataclass
-class bHeader:
-    key: str
-    value: str
-
-@dataclass
-class bRequest:
-    server_ip_address: str
-
-    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: List[bHeader]
-
-    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"]
-        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: int
-    http_version: str
-    reason: str
-    content: str
-    timestamp_start: float
-    timestamp_end: float
-    # [("Header","Data")]
-    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: str
-    request: bRequest
-    response: bResponse
-
-    def __init__(self, flow: dict):
-        self.uid = flow["id"]
-        self.request =  bRequest(flow["request"])
-        self.response =  bResponse(flow["response"])
-
-
-@dataclass
-class bFlowState(Enum):
-    ERROR = 0
-    UNSENT_HTTP_REQUEST = 1
-    SENT_HTTP_REQUEST = 2
-    UNSENT_HTTP_RESPONSE = 3
-    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:
-    state: bFlowState
-    flow: Flow
-    time: float = 0
-    retries_left: int = 5
-
-"""
-
-The network thread communicates with the bigsnitch plugin using zeromq.
-
-"""
-
-class NetworkThread(threading.Thread):
-        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()
-                # 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
-                self.retries = 5
-
-        # 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, flowitem = self.q.get(block=False)
-                    if self.flows.get(i, None):
-                        print(f"flow {i} doubled? ignoring...")
-                        continue
-                    else:
-                        self.flows.append(flowitem)
-
-                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 id, flow in self.flows.items():
-                # send the request if not sent
-                if state == bFlowState.UNSENT_HTTP_REQUEST:
-                    self.send_request(id, bRequest(flow.flow))
-
-                elif state == bFlowState.SENT_HTTP_REQUEST:
-                    # check timer, try resend
-                    delta = time.monotonic() - flows[id].time
-                    if delta > 5:
-                        self.send_request(id, bRequest(flow.flow))
-                        flows[id].retries -= 1
-                        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()
-                try:
-                    if msg:
-                        result = json.loads(msg)
-                        # packet statemachine
-                        if result["msg"] == "ack":
-                            print("m ack received")
-                            return result
-                        else:
-                            print("got unexpected message {result}")
-                except json.JSONDecodeError:
-                    print(f"malformed message received {msg}")
-
-        def run(self):
-                print("thread started")
-                self.connect()
-                while True:
-                    self.timer = time.monotonic()
-                    self.get_new_flows()
-                    self.handle_packets()
-                    self.update_flows()
-
-                    if self.timer - time.monotonic() < -5:
-                        pass
-                        #self.send_msg_and_ack({"msg": "ping"})
-
-        def disconnect(self):
-                self.socket.setsockopt(zmq.LINGER,0)
-                self.socket.close()
-                print("disconnected")
-                
-        def reconnect(self):
-                self.disconnect()
-                time.sleep(1)
-                self.connect()
-        
-        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"})
-                print("connected")
-
-        """
-        def send_msg_and_ack(self, msg):
-                self.timer = time.monotonic()
-                while True:
-                        #print("m sending")
-                        self.send(msg)
-                        if (self.socket.poll(50) & zmq.POLLIN) != 0:
-                                msg = self.socket.recv()
-                                try:
-                                        if msg:
-                                                result = json.loads(msg)
-                                                if result["msg"] == "ack":
-                                                        print("m ack received")
-                                                        return result
-                                                else:
-                                                        print("got unexpected message {result}")
-                                        
-                                except json.JSONDecodeError:
-                                        print(f"malformed message received {msg}")
-                        print("no ack received, reconnecting...")
-                        self.reconnect()
-                return NO_MSG
-        """
 
 class BigSnitchBridge:
     def __init__(self):
+        print("BigSnitchBridge started")
         self.q = Queue()
         self.thread = NetworkThread("network", self.q)
         self.thread.start()

mitmaddon/networkthread.py

@@ -0,0 +1,325 @@
+import pdb
+
+from mitmproxy import ctx
+from mitmproxy.flow import Flow
+import threading
+import time
+import zmq
+import json
+import os
+from enum import Enum
+from dataclasses import dataclass
+from typing import List, Dict
+from queue import Queue, Empty
+
+"""
+# this method is used to convert flow states (generated with get_state()) to json
+def convert_to_strings(obj):
+        if isinstance(obj, dict):
+            return {convert_to_strings(key): convert_to_strings(value)
+                    for key, value in obj.items()}
+        elif isinstance(obj, list) or isinstance(obj, tuple):
+            return [convert_to_strings(element) for element in obj]
+        elif isinstance(obj, bytes):
+            try:
+                data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
+            except:
+                print(obj)
+                data = str(obj)[2:-1]
+            return data
+
+        return obj
+"""
+
+@dataclass
+class bHeader:
+    key: str
+    value: str
+
+
+@dataclass
+class bRequest:
+    server_ip_address: str
+
+    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: List[bHeader]
+
+    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"]
+        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: int
+    http_version: str
+    reason: str
+    content: str
+    timestamp_start: float
+    timestamp_end: float
+    # [("Header","Data")]
+    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: str
+    request: bRequest
+    response: bResponse
+
+    def __init__(self, flow: dict):
+        self.uid = flow["id"]
+        self.request =  bRequest(flow["request"])
+        self.response =  bResponse(flow["response"])
+
+
+@dataclass
+class bFlowState(Enum):
+    ERROR = 0
+    UNSENT_HTTP_REQUEST = 1
+    SENT_HTTP_REQUEST = 2
+    UNSENT_HTTP_RESPONSE = 3
+    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
+
+    def __init__(self, json: Dict):
+        self.ptype = json["type"]
+        self.flowid = int(json["id"])
+        self.data = json["data"]
+
+
+@dataclass
+class FlowItem:
+    state: bFlowState
+    flow: Flow
+    time: float = 0
+    retries_left: int = 5
+
+"""
+
+The network thread communicates with the bigsnitch plugin using zeromq.
+
+"""
+class NetworkThread(threading.Thread):
+        def __init__(self, name: str, queue: Queue, path: str = None):
+                threading.Thread.__init__(self)
+                self.name = name
+                # path
+                self.path = path
+                if not self.path:
+                    self.path = os.environ.get("BIGSNITCH_PATH", None)
+                if not self.path:
+                    self.path = "tcp://127.0.0.1:12345"
+
+                # queue for communicating with the main mitmproxy thread
+                # contains tuples of (id, FlowItem)
+                self.q = queue
+                # for zmq use
+                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
+                self.retries = 5
+
+        # 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, flowitem = self.q.get(block=False)
+                    if self.flows.get(i, None):
+                        print(f"flow {i} doubled? ignoring...")
+                        continue
+                    else:
+                        self.flows.append(flowitem)
+
+                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 id, flow in self.flows.items():
+                # send the request if not sent
+                if state == bFlowState.UNSENT_HTTP_REQUEST:
+                    self.send_request(id, bRequest(flow.flow))
+
+                elif state == bFlowState.SENT_HTTP_REQUEST:
+                    # check timer, try resend
+                    delta = time.monotonic() - flows[id].time
+                    if delta > 5:
+                        self.send_request(id, bRequest(flow.flow))
+                        flows[id].retries -= 1
+                        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()
+                try:
+                    if msg:
+                        result = json.loads(msg)
+                        pkg = bPacket(json=result)
+                        # flow ACKed
+                        if pkg.ptype == bPacketType.ACK:
+                            continue
+                        # flow killed
+                        elif pkg.ptype == bPacketType.KILL:
+                            continue
+                        else:
+                            print(f"got unexpected message {pkg.ptype}")
+                except json.JSONDecodeError:
+                    print(f"malformed message received {msg}")
+
+        def run(self):
+                print("thread started")
+                self.connect()
+                while True:
+                    self.timer = time.monotonic()
+                    self.get_new_flows()
+                    self.handle_packets()
+                    self.update_flows()
+
+                    if self.timer - time.monotonic() < -5:
+                        pass
+                        #self.send_msg_and_ack({"msg": "ping"})
+
+        def disconnect(self):
+                self.socket.setsockopt(zmq.LINGER,0)
+                self.socket.close()
+                print("disconnected")
+
+        def reconnect(self):
+                print("reconnecting")
+                self.disconnect()
+                time.sleep(1)
+                self.connect()
+
+        def connect(self):
+                self.socket = self.context.socket(zmq.PAIR)
+                self.socket.connect(self.path)
+                #self.send_msg_and_ack({"msg": "ping"})
+                print("connected")
+
+        """
+        def send_msg_and_ack(self, msg):
+                self.timer = time.monotonic()
+                while True:
+                        #print("m sending")
+                        self.send(msg)
+                        if (self.socket.poll(50) & zmq.POLLIN) != 0:
+                                msg = self.socket.recv()
+                                try:
+                                        if msg:
+                                                result = json.loads(msg)
+                                                if result["msg"] == "ack":
+                                                        print("m ack received")
+                                                        return result
+                                                else:
+                                                        print("got unexpected message {result}")
+                                        
+                                except json.JSONDecodeError:
+                                        print(f"malformed message received {msg}")
+                        print("no ack received, reconnecting...")
+                        self.reconnect()
+                return NO_MSG
+        """

mitmaddon/test_bigsnitch.py

@@ -1,5 +1,56 @@
 #!/usr/bin/env python3
 
-import pytest
+import pdb
-from bigsnitch import NetworkThread
 
+import pytest
+from networkthread import NetworkThread
+import os
+import tempfile
+from queue import Queue
+import zmq
+
+class MitmAddonTestServer:
+    def __init__(self, queue, path: str):
+        self.queue = queue
+        self.path = path
+        self.socket = None
+        self.context = zmq.Context()
+        self.connect()
+
+    def connect(self):
+        self.socket = self.context.socket(zmq.PAIR)
+        self.socket.connect(self.path)
+
+    def disconnect(self):
+        self.socket.setsockopt(zmq.LINGER,0)
+        self.socket.close()
+
+    def send_packet(self, pkg: bPacket):
+        msg = {"type": pkg.ptype, "id": pkg.flowid, "data": pkg.data}
+        self.send(msg)
+
+@pytest.fixture
+def client_server():
+    queue = Queue()
+    sock = os.path.join(tempfile.mkdtemp(), "bigsnitchtest.sock")
+    sock = f"ipc://{sock}"
+    client = NetworkThread(name="testthread",queue=queue,path=sock)
+    client.daemon = True
+    client.start()
+    server = MitmAddonTestServer(queue, sock)
+    server.connect()
+    yield client, server
+    client.join(1)
+    server.disconnect()
+
+class TestBigSnitchWrapper:
+    def test_request_convert(self):
+        pass
+"""
+class TestMitmAddon:
+    def test_request(self, client_server):
+        self.client, self.server = client_server
+        # create request
+        flowitem = FlowItem(bFlowState.UNSENT_HTTP_REQUEST, flow)
+        self.q.put_nowait((flow.id, flowitem))
+"""