121 lines
4.5 KiB
Python
121 lines
4.5 KiB
Python
#!/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):
|
|
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):
|
|
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()
|
|
]
|