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

def get_msg(socket):
        
        msg = socket.recv()
        try:
                if msg:
                        return json.loads(msg)
        except json.JSONDecodeError:
                print(f"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()
        connected = False
        a = None
        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()

                        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 Counter:
    def __init__(self):
        self.q = Queue()
        self.thread = threading.Thread(name="NetworkThread", target=networking, args=(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 = [
        Counter()
        ]