Merge pull request #1410 from cortesi/addons

Keep maturing scripts and addons

1 files changed, 13 insertions, 12 deletions

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])