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: scripts: refactor some examples that keep global state

We now have a better way to do this.
---
 examples/iframe_injector.py | 37 +++++++++++++++++++------------------
 1 file changed, 19 insertions(+), 18 deletions(-)

(limited to 'examples/iframe_injector.py')

diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py
index 352c3c24..33d18bbd 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])
-- 
cgit v1.2.3