blob: 9cf6575049ff49a765840085e5272b5f8c9a0b9a (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
import os.path
from mitmproxy import exceptions
from mitmproxy import flowfilter
from mitmproxy import io
from mitmproxy import ctx
class Save:
def __init__(self):
self.stream = None
self.filt = None
self.active_flows = set() # type: Set[flow.Flow]
def start_stream_to_path(self, path, mode, flt):
path = os.path.expanduser(path)
try:
f = open(path, mode)
except IOError as v:
raise exceptions.OptionsError(str(v))
self.stream = io.FilteredFlowWriter(f, flt)
self.active_flows = set()
def configure(self, updated):
# We're already streaming - stop the previous stream and restart
if "save_stream_filter" in updated:
if ctx.options.save_stream_filter:
self.filt = flowfilter.parse(ctx.options.save_stream_filter)
if not self.filt:
raise exceptions.OptionsError(
"Invalid filter specification: %s" % ctx.options.save_stream_filter
)
else:
self.filt = None
if "save_stream_file" in updated:
if self.stream:
self.done()
if ctx.options.save_stream_file:
if ctx.options.save_stream_file.startswith("+"):
path = ctx.options.save_stream_file[1:]
mode = "ab"
else:
path = ctx.options.save_stream_file
mode = "wb"
self.start_stream_to_path(path, mode, self.filt)
def tcp_start(self, flow):
if self.stream:
self.active_flows.add(flow)
def tcp_end(self, flow):
if self.stream:
self.stream.add(flow)
self.active_flows.discard(flow)
def response(self, flow):
if self.stream:
self.stream.add(flow)
self.active_flows.discard(flow)
def request(self, flow):
if self.stream:
self.active_flows.add(flow)
def done(self):
if self.stream:
for flow in self.active_flows:
self.stream.add(flow)
self.active_flows = set([])
self.stream.fo.close()
self.stream = None
|
