1 files changed, 18 insertions, 21 deletions
diff --git a/examples/simple/modify_body_inject_iframe.py b/examples/simple/modify_body_inject_iframe.py
index ab5abf27..595bd9f2 100644
--- a/examples/simple/modify_body_inject_iframe.py
+++ b/examples/simple/modify_body_inject_iframe.py
@@ -1,29 +1,26 @@
-# Usage: mitmdump -s "iframe_injector.py url"
 # (this script works best with --anticache)
-import sys
 from bs4 import BeautifulSoup
+from mitmproxy import ctx, http
 
 
 class Injector:
-    def __init__(self, iframe_url):
-        self.iframe_url = iframe_url
+    def load(self, loader):
+        loader.add_option(
+            "iframe", str, "", "IFrame to inject"
+        )
 
-    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")
+    def response(self, flow: http.HTTPFlow) -> None:
+        if ctx.options.iframe:
+            html = BeautifulSoup(flow.response.content, "html.parser")
+            if html.body:
+                iframe = html.new_tag(
+                    "iframe",
+                    src=ctx.options.iframe,
+                    frameborder=0,
+                    height=0,
+                    width=0)
+                html.body.insert(0, iframe)
+                flow.response.content = str(html).encode("utf8")
 
 
-def start(opts):
-    if len(sys.argv) != 2:
-        raise ValueError('Usage: -s "iframe_injector.py url"')
-    return Injector(sys.argv[1])
+addons = [Injector()]