mimimitm

みみ-Manipulator in the Middle
git clone https://git.sinitax.com/sinitax/mimimitm
Log | Files | Refs | README | sfeed.txt

mimimitm.py (3150B)

      1from io import BytesIO
      2from logging import info, warn
      3from typing_extensions import override
      4from mitmproxy import ctx
      5from os import listdir
      6from os.path import basename
      7from PIL import Image, ImageDraw
      8from ultralytics import YOLO
      9from ultralytics.yolo.utils import set_logging
     10from urllib.parse import urlparse
     11
     12import pillow_avif
     13import random
     14
     15set_logging("ultralytics", False)
     16
     17class MIMIMITM:
     18    def __init__(self):
     19        self.model = None
     20
     21    def load(self, loader):
     22        loader.add_option("weights", str, "weights.pt", "weights path (.pt)")
     23        loader.add_option("debug", bool, False, "enable debug")
     24        loader.add_option("imgsize", int, 640, "inference img size")
     25        loader.add_option("conf", float, 0.5, "confidence threshold")
     26        loader.add_option("iou", float, 0.5, "IOU threshold")
     27        loader.add_option("dev", str, "cpu", "inference device")
     28        loader.add_option("augment", bool, True, "inference device")
     29
     30    def running(self):
     31        try:
     32            self.model = YOLO(ctx.options.weights)
     33            self.overlays = [Image.open(f"overlays/{name}") \
     34                for name in listdir("overlays")]
     35            info("MIMIMITM is running!")
     36        except:
     37            addons.clear()
     38            raise
     39
     40    def response(self, flow):
     41        content_type = flow.response.headers.get("content-type", "")
     42
     43        info(f"TYPE {content_type}")
     44
     45        if "image/" not in content_type:
     46            return
     47
     48        if ctx.options.debug:
     49            warn("RUNNING")
     50
     51        try:
     52            io = BytesIO(flow.response.content)
     53            io.name = basename(urlparse(flow.request.url).path)
     54            img = Image.open(io, formats=None)
     55        except:
     56            raise Exception(f"load img of content type {content_type}")
     57
     58        results = self.model.predict(source=img, imgsz=ctx.options.imgsize,
     59            conf=ctx.options.conf, iou=ctx.options.iou,
     60            augment=ctx.options.augment, device=ctx.options.dev)
     61        result = results[0].cpu().numpy()
     62
     63        draw = ImageDraw.Draw(img)
     64        if ctx.options.debug:
     65            warn(f"FOUND {len(result.boxes)}")
     66        for box in result.boxes:
     67            x1, y1, x2, y2 = [int(v + 0.5) for v in box.xyxy[0]]
     68            if ctx.options.debug:
     69                conf = box.conf[0]
     70                info(f"BOX ({conf}) {x1} {y1} {x2} {y2}")
     71                draw.line((x1, y1, x2, y1), fill=(255, 0, 0), width=2)
     72                draw.line((x2, y1, x2, y2), fill=(255, 0, 0), width=2)
     73                draw.line((x2, y2, x1, y2), fill=(255, 0, 0), width=2)
     74                draw.line((x1, y2, x1, y1), fill=(255, 0, 0), width=2)
     75            w, h = abs(x1 - x2), abs(y1 - y2)
     76            if w < h / 2: # facing away approximator
     77                continue
     78            y1, y2 = (y1 - h / 2), (y2 - h / 2)
     79            bx, by = (int(x1 - w / 2), int(y1 - h / 2))
     80            overlay = random.choice(self.overlays).resize((w * 2, h * 2))
     81            img.paste(overlay, box=(bx, by), mask=overlay.split()[-1])
     82
     83        io = BytesIO()
     84        img.save(io, format=img.format)
     85        flow.response.content = io.getvalue()
     86
     87addons = [
     88    MIMIMITM()
     89]