From dbafe9f87bb7b793ae2d84e01af5d39f034c78d4 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 23 Jul 2016 11:57:31 +1200 Subject: scripts: refactor some examples that keep global state We now have a better way to do this. --- examples/flowwriter.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) (limited to 'examples/flowwriter.py') diff --git a/examples/flowwriter.py b/examples/flowwriter.py index 07c7ca20..df2e5a40 100644 --- a/examples/flowwriter.py +++ b/examples/flowwriter.py @@ -3,20 +3,21 @@ import sys from mitmproxy.flow import FlowWriter -state = {} + +class Writer: + def __init__(self, path): + if path == "-": + f = sys.stdout + else: + f = open(path, "wb") + self.w = FlowWriter(f) + + def response(self, flow): + if random.choice([True, False]): + self.w.add(flow) def start(): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') - - if sys.argv[1] == "-": - f = sys.stdout - else: - f = open(sys.argv[1], "wb") - state["flow_writer"] = FlowWriter(f) - - -def response(flow): - if random.choice([True, False]): - state["flow_writer"].add(flow) + return Writer(sys.argv[1]) -- cgit v1.2.3