stuff

mitmaddon/littlesnitch.py

@@ -5,28 +5,75 @@ import threading
 from queue import Queue
 import time
 import zmq
+import json
 
+NO_MSG {"msg": None}
+INIT_MSG = {"msg": "init"}
+ACK_MSG = {"msg": "ack"}
+PING_MSG = {"msg": "ping"}
+PONG_MSG = {"msg": "pong"}
+
+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):
+            return str(obj)[2:-1]
+        return obj
+
+def get_msg(socket):
+        
+        msg = socket.recv()
+        try:
+                if msg:
+                        return json.loads(msg)
+        except json.JSONDecodeError:
+                print("malformed message received '{msg}'")
+
+        return NO_MSG
+
+        
+def send_msg(msg, socket):
+        a = convert_to_strings(msg)
+        socket.send(str.encode(json.dumps(a)))
+        
 def networking(q):
         print("starting thread")
 
         context = zmq.Context()
-        socket = context.socket(zmq.PAIR)
-        socket.connect("tcp://127.0.0.1:12345")
-        
-        while True:
-            print("try recv")
-            message = socket.recv()
-            if message == b"littlesnitch_init":
-                print("connected")
-                socket.send(b"mitmaddon")
-                while True:
-                    a = q.get()
-                    print(f"got {a}")
-                    if a:
-                        socket.send(str.encode(a))
-                    q.task_done()
-            else:
-                raise ValueError("init failed")
+        connected = False
+        while not connected:
+                socket = context.socket(zmq.PAIR)
+                socket.connect("tcp://127.0.0.1:12345")
+                msg = get_msg(socket)
+                if msg["msg"] == "init":
+                        send_msg(ACK_MSG, socket)
+                        connected = True
+
+                timer = time.monotonic()
+                while connected:
+                        if timer - time.monotonic() >= 5:
+                                timer = time.monotonic()
+                                send_msg(PING_MSG,socket)
+                                msg = get_msg(socket)
+                                if msg["msg"] != "pong":
+                                        connected = False
+                        
+                        msg = get_msg(socket)
+                        if msg['msg'] == "ping":
+                                send_msg(PONG_MSG, socket)
+                                timer = time.monotonic()
+                        
+                        a = q.get(block=False)
+                        if a:
+                                msg = get_msg(socket)
+                                if msg["msg"] == "ack":
+                                        q.task_done()
+                                        timer = time.monotonic()
+                                else:
+                                        connected = False
 
 class Counter:
     def __init__(self):
@@ -36,13 +83,11 @@ class Counter:
         self.q.join()
 
     def request(self, flow):
-        data = flow.request.data
-        self.q.put(f"{flow.id},REQ,{data.method},{data.scheme},{data.host},{data.port},{data.path},{data.http_version},{data.headers}")
+        self.q.put(flow.get_state())
         self.q.join()
 
     def response(self, flow):
-        data = flow.response.data
-        self.q.put(f"{flow.id},RES,{data.status_code},{data.http_version},{data.reason},{data.headers},{data.content},{data.timestamp_start},{data.timestamp_end}")
+        self.q.put(flow.get_state())
         self.q.join()