dataclasses, further development
This commit is contained in:
parent
3dbd506427
commit
8a784f9c48
1 changed files with 100 additions and 50 deletions
|
@ -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:
|
||||
|
|
Loading…
Reference in a new issue