This commit is contained in:
Tim Blume 2021-04-25 21:38:58 +02:00
parent 9591c5bebc
commit 0214db88c3
3 changed files with 381 additions and 305 deletions

View file

@ -3,312 +3,12 @@
import pdb
from mitmproxy import ctx
from mitmproxy.flow import Flow
import threading
from queue import Queue, Empty
import time
import zmq
import json
from enum import Enum
from dataclasses import dataclass
from typing import List, Dict
# this method is used to convert flow states (generated with get_state()) to json
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):
try:
data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
except:
print(obj)
data = str(obj)[2:-1]
return data
return obj
#
# bigsnitch communication types
#
@dataclass
class bHeader:
key: str
value: str
@dataclass
class bRequest:
server_ip_address: str
tls: str
content: str
scheme: str
method: str
host: str
port: int
http_version: str
path: str
timestamp_start: float
timestamp_end: float
# [("Header","Data")]
headers: List[bHeader]
error: str
# init from dict
def __init__(self, flow: dict):
self.server_ip_address = flow["server_ip_address"]
self.tls = flow["server_conn"]["tls_established"]
self.content = flow["content"]
self.scheme = flow["scheme"]
self.method = flow["method"]
self.host = flow["host"]
self.port = flow["port"]
self.http_version = flow["http_version"]
self.timestamp_start = flow["timestamp_start"]
self.timestamp_end = flow["timestamp_end"]
for k,v in flow["headers"]:
self.headers.append(bHeader(k,v))
@dataclass
class bResponse:
status_code: int
http_version: str
reason: str
content: str
timestamp_start: float
timestamp_end: float
# [("Header","Data")]
headers: List[bHeader]
def __init__(self, flow: dict):
self.status_code = flow["status_code"]
self.http_version = flow["http_version"]
self.reason = flow["reason"]
self.content = flow["content"]
self.timestamp_start = flow["timestamp_start"]
self.timestamp_end = flow["timestamp_end"]
for k,v in flow["headers"]:
self.headers.append(bHeader(k,v))
@dataclass
class bFlow:
uid: str
request: bRequest
response: bResponse
def __init__(self, flow: dict):
self.uid = flow["id"]
self.request = bRequest(flow["request"])
self.response = bResponse(flow["response"])
@dataclass
class bFlowState(Enum):
ERROR = 0
UNSENT_HTTP_REQUEST = 1
SENT_HTTP_REQUEST = 2
UNSENT_HTTP_RESPONSE = 3
SENT_HTTP_RESPONSE = 4
@dataclass
class bPacketType:
NACK = 0
ACK = 1
KILL = 2
WARNING = 3
ERROR = 4
PING = 5
HTTP_REQUEST = 6
HTTP_RESPONSE = 7
@dataclass
class bPacket:
ptype: bPacketType
flowid: int
data: str
# for use in NetworkThread queue
@dataclass
class FlowItem:
state: bFlowState
flow: Flow
time: float = 0
retries_left: int = 5
"""
The network thread communicates with the bigsnitch plugin using zeromq.
"""
class NetworkThread(threading.Thread):
def __init__(self, name: str, queue: Queue):
threading.Thread.__init__(self)
self.name = name
# queue for communicating with the main mitmproxy thread
# contains tuples of (id, FlowItem)
self.q = queue
# for zmq use
self.context = zmq.Context()
# all current flows being handled by mitmproxy
self.flows: Dict[FlowItem]
# timer for sending pings to check if the connection broke
self.timer = time.monotonic()
# retries left for reconnecting / resending a broken flow
self.retries = 5
# send a single message, no checks involved
def send(self, msg):
a = convert_to_strings(msg)
self.socket.send(str.encode(json.dumps(a)))
# add new flows from the queue
def get_new_flows(self):
while True:
try:
# get new flows that may occured
i, flowitem = self.q.get(block=False)
if self.flows.get(i, None):
print(f"flow {i} doubled? ignoring...")
continue
else:
self.flows.append(flowitem)
except Empty:
break
def send_packet(self, pkg: bPacket):
msg = {"type": pkg.ptype, "id": pkg.flowid, "data": pkg.data}
self.send(msg)
def send_http_request(self, id: int, request: bRequest):
pkg = bPacket(bPacketType.HTTP_REQUEST, id, request.json())
self.send_packet(pkg)
flows[id].state = bFlowState.SENT_HTTP_REQUEST
flows[id].time = time.monotonic()
def send_http_response(self, id: int, response: bResponse):
pkg = bPacket(bPacketType.HTTP_RESPONSE, id, response.json())
self.send_packet(pkg)
flows[id].state = bFlowState.SENT_HTTP_RESPONSE
flows[id].time = time.monotonic()
# update all current flows
# handles the state machine for each flow
def update_flows(self):
for id, flow in self.flows.items():
# send the request if not sent
if state == bFlowState.UNSENT_HTTP_REQUEST:
self.send_request(id, bRequest(flow.flow))
elif state == bFlowState.SENT_HTTP_REQUEST:
# check timer, try resend
delta = time.monotonic() - flows[id].time
if delta > 5:
self.send_request(id, bRequest(flow.flow))
flows[id].retries -= 1
continue
if flows[id].retries <= 0:
flows[id].flow.kill()
print(f"http request {id} timed out! flow killed.")
# send the response if not sent
elif state == bFlowState.UNSENT_HTTP_RESPONSE:
self.send_response(id, bResponse(flow.flow))
elif state == bFlowState.SENT_HTTP_RESPONSE:
# check timer, try resend
delta = time.monotonic() - flows[id].time
if delta > 5:
self.send_response(id, bResponse(flow.flow))
flows[id].retries -= 1
continue
if flows[id].retries <= 0:
flows[id].flow.kill()
print(f"http response {id} timedout! flow killed.")
elif state == bFlowState.ERROR:
print(f"error in flow {id}!")
# handle incoming packets / update the statemachine
def handle_packets(self):
while((self.socket.poll(50) & zmq.POLLIN) != 0):
msg = self.socket.recv()
try:
if msg:
result = json.loads(msg)
# packet statemachine
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}")
def run(self):
print("thread started")
self.connect()
while True:
self.timer = time.monotonic()
self.get_new_flows()
self.handle_packets()
self.update_flows()
if self.timer - time.monotonic() < -5:
pass
#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("connected")
"""
def send_msg_and_ack(self, msg):
self.timer = time.monotonic()
while True:
#print("m sending")
self.send(msg)
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
"""
from networkthread import bFlowState, FlowItem, NetworkThread
class BigSnitchBridge:
def __init__(self):
print("BigSnitchBridge started")
self.q = Queue()
self.thread = NetworkThread("network", self.q)
self.thread.start()

325
mitmaddon/networkthread.py Normal file
View file

@ -0,0 +1,325 @@
import pdb
from mitmproxy import ctx
from mitmproxy.flow import Flow
import threading
import time
import zmq
import json
import os
from enum import Enum
from dataclasses import dataclass
from typing import List, Dict
from queue import Queue, Empty
"""
# this method is used to convert flow states (generated with get_state()) to json
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):
try:
data = obj.decode('unicode-escape').encode('latin1').decode('utf-8')
except:
print(obj)
data = str(obj)[2:-1]
return data
return obj
"""
@dataclass
class bHeader:
key: str
value: str
@dataclass
class bRequest:
server_ip_address: str
tls: str
content: str
scheme: str
method: str
host: str
port: int
http_version: str
path: str
timestamp_start: float
timestamp_end: float
# [("Header","Data")]
headers: List[bHeader]
error: str
# init from dict
def __init__(self, flow: dict):
self.server_ip_address = flow["server_ip_address"]
self.tls = flow["server_conn"]["tls_established"]
self.content = flow["content"]
self.scheme = flow["scheme"]
self.method = flow["method"]
self.host = flow["host"]
self.port = flow["port"]
self.http_version = flow["http_version"]
self.timestamp_start = flow["timestamp_start"]
self.timestamp_end = flow["timestamp_end"]
for k,v in flow["headers"]:
self.headers.append(bHeader(k,v))
@dataclass
class bResponse:
status_code: int
http_version: str
reason: str
content: str
timestamp_start: float
timestamp_end: float
# [("Header","Data")]
headers: List[bHeader]
def __init__(self, flow: dict):
self.status_code = flow["status_code"]
self.http_version = flow["http_version"]
self.reason = flow["reason"]
self.content = flow["content"]
self.timestamp_start = flow["timestamp_start"]
self.timestamp_end = flow["timestamp_end"]
for k,v in flow["headers"]:
self.headers.append(bHeader(k,v))
@dataclass
class bFlow:
uid: str
request: bRequest
response: bResponse
def __init__(self, flow: dict):
self.uid = flow["id"]
self.request = bRequest(flow["request"])
self.response = bResponse(flow["response"])
@dataclass
class bFlowState(Enum):
ERROR = 0
UNSENT_HTTP_REQUEST = 1
SENT_HTTP_REQUEST = 2
UNSENT_HTTP_RESPONSE = 3
SENT_HTTP_RESPONSE = 4
@dataclass
class bPacketType:
NACK = 0
ACK = 1
KILL = 2
WARNING = 3
ERROR = 4
PING = 5
HTTP_REQUEST = 6
HTTP_RESPONSE = 7
@dataclass
class bPacket:
ptype: bPacketType
flowid: int
data: str
def __init__(self, json: Dict):
self.ptype = json["type"]
self.flowid = int(json["id"])
self.data = json["data"]
@dataclass
class FlowItem:
state: bFlowState
flow: Flow
time: float = 0
retries_left: int = 5
"""
The network thread communicates with the bigsnitch plugin using zeromq.
"""
class NetworkThread(threading.Thread):
def __init__(self, name: str, queue: Queue, path: str = None):
threading.Thread.__init__(self)
self.name = name
# path
self.path = path
if not self.path:
self.path = os.environ.get("BIGSNITCH_PATH", None)
if not self.path:
self.path = "tcp://127.0.0.1:12345"
# queue for communicating with the main mitmproxy thread
# contains tuples of (id, FlowItem)
self.q = queue
# for zmq use
self.context = zmq.Context()
# all current flows being handled by mitmproxy
self.flows: Dict[FlowItem] = {}
# timer for sending pings to check if the connection broke
self.timer = time.monotonic()
# retries left for reconnecting
self.retries = 5
# send a single message, no checks involved
def send(self, msg):
a = convert_to_strings(msg)
self.socket.send(str.encode(json.dumps(a)))
# add new flows from the queue
def get_new_flows(self):
while True:
try:
# get new flows that may occured
i, flowitem = self.q.get(block=False)
if self.flows.get(i, None):
print(f"flow {i} doubled? ignoring...")
continue
else:
self.flows.append(flowitem)
except Empty:
break
def send_packet(self, pkg: bPacket):
msg = {"type": pkg.ptype, "id": pkg.flowid, "data": pkg.data}
self.send(msg)
def send_http_request(self, id: int, request: bRequest):
pkg = bPacket(bPacketType.HTTP_REQUEST, id, request.json())
self.send_packet(pkg)
flows[id].state = bFlowState.SENT_HTTP_REQUEST
flows[id].time = time.monotonic()
def send_http_response(self, id: int, response: bResponse):
pkg = bPacket(bPacketType.HTTP_RESPONSE, id, response.json())
self.send_packet(pkg)
flows[id].state = bFlowState.SENT_HTTP_RESPONSE
flows[id].time = time.monotonic()
# update all current flows
# handles the state machine for each flow
def update_flows(self):
for id, flow in self.flows.items():
# send the request if not sent
if state == bFlowState.UNSENT_HTTP_REQUEST:
self.send_request(id, bRequest(flow.flow))
elif state == bFlowState.SENT_HTTP_REQUEST:
# check timer, try resend
delta = time.monotonic() - flows[id].time
if delta > 5:
self.send_request(id, bRequest(flow.flow))
flows[id].retries -= 1
continue
if flows[id].retries <= 0:
flows[id].flow.kill()
print(f"http request {id} timed out! flow killed.")
# send the response if not sent
elif state == bFlowState.UNSENT_HTTP_RESPONSE:
self.send_response(id, bResponse(flow.flow))
elif state == bFlowState.SENT_HTTP_RESPONSE:
# check timer, try resend
delta = time.monotonic() - flows[id].time
if delta > 5:
self.send_response(id, bResponse(flow.flow))
flows[id].retries -= 1
continue
if flows[id].retries <= 0:
flows[id].flow.kill()
print(f"http response {id} timedout! flow killed.")
elif state == bFlowState.ERROR:
print(f"error in flow {id}!")
# handle incoming packets / update the statemachine
def handle_packets(self):
while((self.socket.poll(50) & zmq.POLLIN) != 0):
msg = self.socket.recv()
try:
if msg:
result = json.loads(msg)
pkg = bPacket(json=result)
# flow ACKed
if pkg.ptype == bPacketType.ACK:
continue
# flow killed
elif pkg.ptype == bPacketType.KILL:
continue
else:
print(f"got unexpected message {pkg.ptype}")
except json.JSONDecodeError:
print(f"malformed message received {msg}")
def run(self):
print("thread started")
self.connect()
while True:
self.timer = time.monotonic()
self.get_new_flows()
self.handle_packets()
self.update_flows()
if self.timer - time.monotonic() < -5:
pass
#self.send_msg_and_ack({"msg": "ping"})
def disconnect(self):
self.socket.setsockopt(zmq.LINGER,0)
self.socket.close()
print("disconnected")
def reconnect(self):
print("reconnecting")
self.disconnect()
time.sleep(1)
self.connect()
def connect(self):
self.socket = self.context.socket(zmq.PAIR)
self.socket.connect(self.path)
#self.send_msg_and_ack({"msg": "ping"})
print("connected")
"""
def send_msg_and_ack(self, msg):
self.timer = time.monotonic()
while True:
#print("m sending")
self.send(msg)
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
"""

View file

@ -1,5 +1,56 @@
#!/usr/bin/env python3
import pytest
from bigsnitch import NetworkThread
import pdb
import pytest
from networkthread import NetworkThread
import os
import tempfile
from queue import Queue
import zmq
class MitmAddonTestServer:
def __init__(self, queue, path: str):
self.queue = queue
self.path = path
self.socket = None
self.context = zmq.Context()
self.connect()
def connect(self):
self.socket = self.context.socket(zmq.PAIR)
self.socket.connect(self.path)
def disconnect(self):
self.socket.setsockopt(zmq.LINGER,0)
self.socket.close()
def send_packet(self, pkg: bPacket):
msg = {"type": pkg.ptype, "id": pkg.flowid, "data": pkg.data}
self.send(msg)
@pytest.fixture
def client_server():
queue = Queue()
sock = os.path.join(tempfile.mkdtemp(), "bigsnitchtest.sock")
sock = f"ipc://{sock}"
client = NetworkThread(name="testthread",queue=queue,path=sock)
client.daemon = True
client.start()
server = MitmAddonTestServer(queue, sock)
server.connect()
yield client, server
client.join(1)
server.disconnect()
class TestBigSnitchWrapper:
def test_request_convert(self):
pass
"""
class TestMitmAddon:
def test_request(self, client_server):
self.client, self.server = client_server
# create request
flowitem = FlowItem(bFlowState.UNSENT_HTTP_REQUEST, flow)
self.q.put_nowait((flow.id, flowitem))
"""