Merge pull request #1410 from cortesi/addons

Keep maturing scripts and addons

1 files changed, 19 insertions, 18 deletions

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