1 files changed, 22 insertions, 20 deletions
diff --git a/examples/simple/modify_body_inject_iframe.py b/examples/simple/modify_body_inject_iframe.py
index 442a5118..d54468d2 100644
--- a/examples/simple/modify_body_inject_iframe.py
+++ b/examples/simple/modify_body_inject_iframe.py
@@ -1,29 +1,31 @@
-# Usage: mitmdump -s "iframe_injector.py url"
 # (this script works best with --anticache)
-import sys
 from bs4 import BeautifulSoup
 
 
 class Injector:
-    def __init__(self, iframe_url):
-        self.iframe_url = iframe_url
+    def __init__(self):
+        self.iframe_url = None
+
+    def load(self, loader):
+        loader.add_option(
+            "iframe", str, "", "IFrame to inject"
+        )
+
+    def configure(self, options, updated):
+        self.iframe_url = options.iframe
 
     def response(self, flow):
-        if flow.request.host in self.iframe_url:
-            return
-        html = BeautifulSoup(flow.response.content, "html.parser")
-        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")
+        if self.iframe_url:
+            html = BeautifulSoup(flow.response.content, "html.parser")
+            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 load(l):
-    if len(sys.argv) != 2:
-        raise ValueError('Usage: -s "iframe_injector.py url"')
-    return l.boot_into(Injector(sys.argv[1]))
+addons = [Injector()]