From dbafe9f87bb7b793ae2d84e01af5d39f034c78d4 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sat, 23 Jul 2016 11:57:31 +1200
Subject: [PATCH] scripts: refactor some examples that keep global state

We now have a better way to do this.
---
 examples/filt.py            | 20 +++++++++++---------
 examples/flowwriter.py      | 25 +++++++++++++------------
 examples/iframe_injector.py | 37 +++++++++++++++++++------------------
 3 files changed, 43 insertions(+), 39 deletions(-)

diff --git a/examples/filt.py b/examples/filt.py
index 21744edd8..9ccf9fa12 100644
--- a/examples/filt.py
+++ b/examples/filt.py
@@ -1,18 +1,20 @@
-# This scripts demonstrates how to use mitmproxy's filter pattern in inline scripts.
+# This scripts demonstrates how to use mitmproxy's filter pattern in scripts.
 # Usage: mitmdump -s "filt.py FILTER"
 import sys
 from mitmproxy import filt
 
-state = {}
+
+class Filter:
+    def __init__(self, spec):
+        self.filter = filt.parse(spec)
+
+    def response(self, flow):
+        if flow.match(self.filter):
+            print("Flow matches filter:")
+            print(flow)
 
 
 def start():
     if len(sys.argv) != 2:
         raise ValueError("Usage: -s 'filt.py FILTER'")
-    state["filter"] = filt.parse(sys.argv[1])
-
-
-def response(flow):
-    if flow.match(state["filter"]):
-        print("Flow matches filter:")
-        print(flow)
+    return Filter(sys.argv[1])
diff --git a/examples/flowwriter.py b/examples/flowwriter.py
index 07c7ca200..df2e5a40f 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])
diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py
index 352c3c24a..33d18bbd2 100644
--- a/examples/iframe_injector.py
+++ b/examples/iframe_injector.py
@@ -3,26 +3,27 @@
 import sys
 from bs4 import BeautifulSoup
 
-iframe_url = None
+
+class Injector:
+    def __init__(self, iframe_url):
+        self.iframe_url = iframe_url
+
+    def response(self, flow):
+        if flow.request.host in self.iframe_url:
+            return
+        html = BeautifulSoup(flow.response.content, "lxml")
+        if html.body:
+            iframe = html.new_tag(
+                "iframe",
+                src=self.iframe_url,
+                frameborder=0,
+                height=0,
+                width=0)
+            html.body.insert(0, iframe)
+            flow.response.content = str(html).encode("utf8")
 
 
 def start():
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "iframe_injector.py url"')
-    global iframe_url
-    iframe_url = sys.argv[1]
-
-
-def response(flow):
-    if flow.request.host in iframe_url:
-        return
-    html = BeautifulSoup(flow.response.content, "lxml")
-    if html.body:
-        iframe = html.new_tag(
-            "iframe",
-            src=iframe_url,
-            frameborder=0,
-            height=0,
-            width=0)
-        html.body.insert(0, iframe)
-        flow.response.content = str(html).encode("utf8")
+    return Injector(sys.argv[1])