Make a silly little thingHEADmaster

16 files changed, 112 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..9065b4f
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__
+*.pt
diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..d1b2f44
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,10 @@
+all: run
+
+weights.pt:
+	wget "https://drive.google.com/u/0/uc?id=1qcr9DbgsX3ryrz2uU8w4Xm3cOrRywXqb&export=download" -O $@
+
+venv:
+	python3 -m virtualenv venv && source venv/bin/activate && python3 -m pip install -r requirements.txt
+
+run: | venv weights.pt
+	source venv/bin/activate && mitmproxy -s mimimitm.py
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..3ba773e
--- /dev/null
+++ b/README.md
@@ -0,0 +1,8 @@
+みみ-Manipulator in the Middle
+==============================
+
+Based on a [tweet by SwiftOnSecurity](https://twitter.com/SwiftOnSecurity/status/1674137459507056648).
+
+![Browser Example](examples/browser.png)
+
+Remember to install the mitmproxy root certificate for HTTPS support.
diff --git a/examples/browser.png b/examples/browser.png
new file mode 100644
index 0000000..f198ea3
--- /dev/null
+++ b/examples/browser.png
diff --git a/examples/collage.jpg b/examples/collage.jpg
new file mode 100644
index 0000000..980691d
--- /dev/null
+++ b/examples/collage.jpg
diff --git a/mimimitm.py b/mimimitm.py
new file mode 100644
index 0000000..a1e620b
--- /dev/null
+++ b/mimimitm.py
@@ -0,0 +1,89 @@
+from io import BytesIO
+from logging import info, warn
+from typing_extensions import override
+from mitmproxy import ctx
+from os import listdir
+from os.path import basename
+from PIL import Image, ImageDraw
+from ultralytics import YOLO
+from ultralytics.yolo.utils import set_logging
+from urllib.parse import urlparse
+
+import pillow_avif
+import random
+
+set_logging("ultralytics", False)
+
+class MIMIMITM:
+    def __init__(self):
+        self.model = None
+
+    def load(self, loader):
+        loader.add_option("weights", str, "weights.pt", "weights path (.pt)")
+        loader.add_option("debug", bool, False, "enable debug")
+        loader.add_option("imgsize", int, 640, "inference img size")
+        loader.add_option("conf", float, 0.5, "confidence threshold")
+        loader.add_option("iou", float, 0.5, "IOU threshold")
+        loader.add_option("dev", str, "cpu", "inference device")
+        loader.add_option("augment", bool, True, "inference device")
+
+    def running(self):
+        try:
+            self.model = YOLO(ctx.options.weights)
+            self.overlays = [Image.open(f"overlays/{name}") \
+                for name in listdir("overlays")]
+            info("MIMIMITM is running!")
+        except:
+            addons.clear()
+            raise
+
+    def response(self, flow):
+        content_type = flow.response.headers.get("content-type", "")
+
+        info(f"TYPE {content_type}")
+
+        if "image/" not in content_type:
+            return
+
+        if ctx.options.debug:
+            warn("RUNNING")
+
+        try:
+            io = BytesIO(flow.response.content)
+            io.name = basename(urlparse(flow.request.url).path)
+            img = Image.open(io, formats=None)
+        except:
+            raise Exception(f"load img of content type {content_type}")
+
+        results = self.model.predict(source=img, imgsz=ctx.options.imgsize,
+            conf=ctx.options.conf, iou=ctx.options.iou,
+            augment=ctx.options.augment, device=ctx.options.dev)
+        result = results[0].cpu().numpy()
+
+        draw = ImageDraw.Draw(img)
+        if ctx.options.debug:
+            warn(f"FOUND {len(result.boxes)}")
+        for box in result.boxes:
+            x1, y1, x2, y2 = [int(v + 0.5) for v in box.xyxy[0]]
+            if ctx.options.debug:
+                conf = box.conf[0]
+                info(f"BOX ({conf}) {x1} {y1} {x2} {y2}")
+                draw.line((x1, y1, x2, y1), fill=(255, 0, 0), width=2)
+                draw.line((x2, y1, x2, y2), fill=(255, 0, 0), width=2)
+                draw.line((x2, y2, x1, y2), fill=(255, 0, 0), width=2)
+                draw.line((x1, y2, x1, y1), fill=(255, 0, 0), width=2)
+            w, h = abs(x1 - x2), abs(y1 - y2)
+            if w < h / 2: # facing away approximator
+                continue
+            y1, y2 = (y1 - h / 2), (y2 - h / 2)
+            bx, by = (int(x1 - w / 2), int(y1 - h / 2))
+            overlay = random.choice(self.overlays).resize((w * 2, h * 2))
+            img.paste(overlay, box=(bx, by), mask=overlay.split()[-1])
+
+        io = BytesIO()
+        img.save(io, format=img.format)
+        flow.response.content = io.getvalue()
+
+addons = [
+    MIMIMITM()
+]
diff --git a/overlays/1.png b/overlays/1.png
new file mode 100644
index 0000000..889d18c
--- /dev/null
+++ b/overlays/1.png
diff --git a/overlays/2.png b/overlays/2.png
new file mode 100644
index 0000000..4f4ce13
--- /dev/null
+++ b/overlays/2.png
diff --git a/overlays/3.png b/overlays/3.png
new file mode 100644
index 0000000..edf578a
--- /dev/null
+++ b/overlays/3.png
diff --git a/overlays/4.png b/overlays/4.png
new file mode 100644
index 0000000..10ef74c
--- /dev/null
+++ b/overlays/4.png
diff --git a/overlays/5.png b/overlays/5.png
new file mode 100644
index 0000000..2622503
--- /dev/null
+++ b/overlays/5.png
diff --git a/overlays/6.png b/overlays/6.png
new file mode 100644
index 0000000..b9b5b67
--- /dev/null
+++ b/overlays/6.png
diff --git a/overlays/7.png b/overlays/7.png
new file mode 100644
index 0000000..e41f5ca
--- /dev/null
+++ b/overlays/7.png
diff --git a/overlays/8.png b/overlays/8.png
new file mode 100644
index 0000000..4b732d7
--- /dev/null
+++ b/overlays/8.png
diff --git a/overlays/9.png b/overlays/9.png
new file mode 100644
index 0000000..ed5082b
--- /dev/null
+++ b/overlays/9.png
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..1ca7a83
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,3 @@
+ultralytics==8.0.124
+pillow-avif-plugin==1.3.1
+mitmproxy==9.0.1