dataclasses, further development

This commit is contained in:
end 2021-04-22 00:41:13 +02:00
parent 3dbd506427
commit 8a784f9c48

View file

@ -7,6 +7,8 @@ import time
import zmq
import json
from enum import Enum
from dataclasses import dataclass
from typing import List
# this method is used to convert flow states (generated with get_state()) to json
def convert_to_strings(obj):
@ -25,46 +27,87 @@ def convert_to_strings(obj):
return data
return obj
#
# bigsnitch communication types
#
# bigsnitch Request type
@dataclass
class bHeader:
key: str
value: str
@dataclass
class bRequest:
server_ip_address = ""
server_ip_address: str
tls = ""
content = ""
scheme = ""
method = ""
host = ""
port = 0
http_version = ""
path = ""
timestamp_start = 0.0
timestamp_end = 0.0
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 = []
headers: List[bHeader]
error = ""
error: str
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 = 0
http_version = ""
reason = ""
content = ""
timestamp_start = 0.0
timestamp_end = 0.0
status_code: int
http_version: str
reason: str
content: str
timestamp_start: float
timestamp_end: float
# [("Header","Data")]
headers = []
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 = "";
request = None
response = None
uid: str;
request: bRequest
response: bResponse
def __init__(self, flow: dict):
self.uid = flow["id"]
self.request = bRequest(flow["request"])
self.response = bRequest(flow["response"])
class FlowState(Enum):
UNSENT_REQ = 0
SENT_REQ = 1
UNSENT_RES = 2
SENT_REQ = 3
SENT_RES = 3
"""
@ -76,40 +119,51 @@ class NetworkThread(threading.Thread):
def __init__(self, name, queue):
threading.Thread.__init__(self)
self.name = name
# queue for communicating with the main mitmproxy thread
self.q = queue
# all current flows being handled by mitmproxy
# id : (state, flow, timer, retries left)
self.flows = {}
self.context = zmq.Context()
# timer for sending pings
# timer for sending pings to check if the connection broke
self.timer = time.monotonic()
# retries left for reconnecting
# retries left for reconnecting / resending a broken flow
self.retries = 5
# get new flows that may occured
def update_flows(self):
# 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, flow, typ = self.q.get(block=False)
if flows.get(i, None):
print("flow {} doubled? ignoring...")
if self.flows.get(i, None):
print(f"flow {i} doubled? ignoring...")
continue
# csave the new flows, if necessary
if typ == "request":
flows[i] = (FlowState.UNSENT_REQ, flow, self.timer, 5)
self.flows[i] = (FlowState.UNSENT_REQ, flow, time.monotonic(), 5)
elif typ == "response":
flows[i] = (FlowState.UNSENT_RES, flow, self.timer, 5)
self.flows[i] = (FlowState.UNSENT_RES, flow, time.monotonic(), 5)
except Empty:
break
# state machine for flows
def handle_flow(self, state, flow):
for i, v in flows.items:
# update all current flows
def update_flows(self):
for k,v in self.flows.items():
state, flow, timer, retries = v
# state machine for flows
if state == FlowState.UNSENT_REQ:
# send it
a = convert_to_strings(msg)
self.socket.send(str.encode(json.dumps(a)))
msg = b""
# send the request
self.send(msg)
pass
elif state == FlowState.SENT_REQ:
# check timer, try resend
@ -119,8 +173,6 @@ class NetworkThread(threading.Thread):
elif state == FlowState.SENT_RES:
pass
self.send_msg_and_ack(a)
def handle_packets(self):
while((self.socket.poll(50) & zmq.POLLIN) != 0):
msg = self.socket.recv()
@ -141,14 +193,13 @@ class NetworkThread(threading.Thread):
self.connect()
while True:
self.timer = time.monotonic()
update_flows()
self.get_new_flows()
self.handle_packets()
self.update_flows()
if self.timer - time.monotonic() < -5:
self.send_msg_and_ack({"msg": "ping"})
def send(self, msg):
a = convert_to_strings(msg)
self.socket.send(str.encode(json.dumps(a)))
pass
#self.send_msg_and_ack({"msg": "ping"})
def disconnect(self):
self.socket.setsockopt(zmq.LINGER,0)
@ -169,9 +220,8 @@ class NetworkThread(threading.Thread):
def send_msg_and_ack(self, msg):
self.timer = time.monotonic()
while True:
print("m sending")
a = convert_to_strings(msg)
self.socket.send(str.encode(json.dumps(a)))
#print("m sending")
self.send(msg)
if (self.socket.poll(50) & zmq.POLLIN) != 0:
msg = self.socket.recv()
try: