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
        
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):
            data = ""
            try:
                data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
            except:
                data = str(obj)[2:-1]
            return data

        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):
                print("thread started")
                self.connect()
                while True:
                        timer = time.monotonic()
                        a = None
                        if not a:
                                try:
                                        a = self.q.get(block=False)
                                except Empty:
                                        pass
                        if a:
                                self.send_msg_and_ack(a)
                                timer = time.monotonic()
                        if timer - time.monotonic() < -5:
                                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("successfully connected")
        
        def send_msg_and_ack(self, msg):
                while True:
                        print("m sending")
                        a = convert_to_strings(msg)
                        self.socket.send(str.encode(json.dumps(a)))
                        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 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()})

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

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

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

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

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

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

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

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

    def websocket_end(self, flow):
        self.q.put({'msg': 'websocket_end', 'flow': flow.get_state()})
        
addons = [
        LittleSnitchBridge()
        ]