mirror of
https://github.com/vee1e/mitmproxy.git
synced 2026-09-03 03:07:18 +00:00
24cf8da27e
The primary motivation here (and for all the other moving around) is to present a clean "front of house" to library users, and to migrate primary objects to the top of the module hierarchy.
79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
import os
|
|
|
|
from mitmproxy import exceptions
|
|
from mitmproxy import flowfilter
|
|
from mitmproxy import http
|
|
from mitmproxy import tcp
|
|
from mitmproxy.contrib import tnetstring
|
|
from mitmproxy import io_compat
|
|
|
|
|
|
FLOW_TYPES = dict(
|
|
http=http.HTTPFlow,
|
|
tcp=tcp.TCPFlow,
|
|
)
|
|
|
|
|
|
class FlowWriter:
|
|
def __init__(self, fo):
|
|
self.fo = fo
|
|
|
|
def add(self, flow):
|
|
d = flow.get_state()
|
|
tnetstring.dump(d, self.fo)
|
|
|
|
|
|
class FlowReader:
|
|
def __init__(self, fo):
|
|
self.fo = fo
|
|
|
|
def stream(self):
|
|
"""
|
|
Yields Flow objects from the dump.
|
|
"""
|
|
try:
|
|
while True:
|
|
data = tnetstring.load(self.fo)
|
|
try:
|
|
data = io_compat.migrate_flow(data)
|
|
except ValueError as e:
|
|
raise exceptions.FlowReadException(str(e))
|
|
if data["type"] not in FLOW_TYPES:
|
|
raise exceptions.FlowReadException("Unknown flow type: {}".format(data["type"]))
|
|
yield FLOW_TYPES[data["type"]].from_state(data)
|
|
except ValueError as e:
|
|
if str(e) == "not a tnetstring: empty file":
|
|
return # Error is due to EOF
|
|
raise exceptions.FlowReadException("Invalid data format.")
|
|
|
|
|
|
class FilteredFlowWriter:
|
|
def __init__(self, fo, flt):
|
|
self.fo = fo
|
|
self.flt = flt
|
|
|
|
def add(self, flow):
|
|
if self.flt and not flowfilter.match(self.flt, flow):
|
|
return
|
|
d = flow.get_state()
|
|
tnetstring.dump(d, self.fo)
|
|
|
|
|
|
def read_flows_from_paths(paths):
|
|
"""
|
|
Given a list of filepaths, read all flows and return a list of them.
|
|
From a performance perspective, streaming would be advisable -
|
|
however, if there's an error with one of the files, we want it to be raised immediately.
|
|
|
|
Raises:
|
|
FlowReadException, if any error occurs.
|
|
"""
|
|
try:
|
|
flows = []
|
|
for path in paths:
|
|
path = os.path.expanduser(path)
|
|
with open(path, "rb") as f:
|
|
flows.extend(FlowReader(f).stream())
|
|
except IOError as e:
|
|
raise exceptions.FlowReadException(e.strerror)
|
|
return flows
|