101 lines
3.2 KiB
Python
101 lines
3.2 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
|
|
|
|
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(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()
|
|
]
|