bigsnitch/mitmaddon/littlesnitch.py

#!/usr/bin/env python3
import pdb
from mitmproxy import ctx
import threading
from queue import Queue, Empty
import time
import zmq
import json
from enum import Enum

NO_MSG = {"msg": None}
INIT_MSG = {"msg": "init"}
ACK_MSG = {"msg": "ack"}
PING_MSG = {"msg": "ping"}
PONG_MSG = {"msg": "pong"}

class NetworkState(Enum):
        DISCONNECTED = auto()
        CONNECTED = auto()
        PING = auto()
        SENDING = auto()
        
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

class NetworkThread(threading.Thread):
        def __init__(self, name, queue):
                threading.Thread.__init__(self)
                self.name = name
                self.q = queue
                self.context = zmq.Context()

        def run(self):
                self.connect()
                msg = self.send_msg_and_expect()

        def disconnect(self):
                self.socket.setsockopt(zmq.LINGER,0)
                self.socket.close()
                
        def reconnect(self):
                self.disconnect()
                self.connect()
        
        def connect(self):
                self.socket = self.context.socket(zmq.PAIR)
                self.socket.connect("tcp://127.0.0.1:12345")
        
        def send_msg_and_expect(this, msg, expect, timeout=5, retries=3):
                while retries:
                        a = convert_to_strings(msg)
                        self.socket.send(str.encode(json.dumps(a)))
                        if (client.poll(REQUEST_TIMEOUT) & zmq.POLLIN) != 0:
                                msg = self.socket.recv()
                                try:
                                        if msg:
                                                result = json.loads(msg)
                                                if result["msg"] in expect:
                                                        return result
                                                else:
                                                        print("got unexpected message {result}")
                                        
                                except json.JSONDecodeError:
                                        print(f"malformed message received {msg}")
                        retries -= 1
                        self.reconnect()
                return NO_MSG

        """
def networking(q):
        print("starting thread")

        state = NetworkState.DISCONNECTED
        a = None
        while state == NetworkState.DISCONNECTED:
                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)
                        state = NetworkState.CONNECTED
                        


                        
                timer = time.monotonic()
                while state != NetworkState.DISCONNECTED:
                        if state == NetworkState.CONNECTED and timer - time.monotonic() >= 5:
                                timer = time.monotonic()
                                send_msg(PING_MSG,socket)
                                msg = get_msg(socket)
                                if msg["msg"] != "pong":
                                        state = NetworkState.
                        
                        msg = get_msg(socket)
                        if msg['msg'] == "ping":
                                send_msg(PONG_MSG, socket)
                                timer = time.monotonic()

                        if not a:
                                try:
                                        a = q.get(block=False)
                                except Empty:
                                        pass
                        if a:
                                send_msg(a, socket)
                                msg = get_msg(socket)
                                if msg["msg"] == "ack":
                                        timer = time.monotonic()
                                        a = None
                                        self.q.task_done()
                                else:
                                        connected = False
"""
class LittleSnitchBridge:
    def __init__(self):
        self.q = Queue()
        self.thread = NetworkThread("network", self.q)
        self.thread.start()

    def request(self, flow):
        self.q.put({'msg': 'request', 'flow': flow.get_state()})
        self.q.join()

    def response(self, flow):
        self.q.put({'msg': 'response', 'flow': flow.get_state()})
        self.q.join()

addons = [
        LittleSnitchBridge()
        ]