Driftnet Reborn - Run with `mitmdump -q -s driftnet.py`

driftnet.py

from mitmproxy import ctx
import multiprocessing

# Define the run_gui function before it is used
def run_gui(queue):
    from PyQt6.QtWidgets import (
        QApplication, QLabel, QVBoxLayout, QWidget, QScrollArea, QGridLayout
    )
    from PyQt6.QtGui import QPixmap
    from PyQt6.QtCore import Qt
    import requests
    import sys

    class ImageViewer(QWidget):
        def __init__(self, queue):
            super().__init__()
            self.queue = queue
            self.setWindowTitle("Image Viewer")
            self.layout = QVBoxLayout()
            self.setLayout(self.layout)

            # Add scroll area to handle large numbers of images
            self.scroll_area = QScrollArea(self)
            self.scroll_area.setWidgetResizable(True)
            self.container = QWidget()
            self.grid_layout = QGridLayout(self.container)
            self.scroll_area.setWidget(self.container)
            self.layout.addWidget(self.scroll_area)

            # Grid layout tracking
            self.row = 0
            self.col = 0
            self.columns = 3  # Number of columns for tiled layout

            # Start a timer to check the queue for new URLs
            self.start_timer()

        def start_timer(self):
            from PyQt6.QtCore import QTimer
            self.timer = QTimer(self)
            self.timer.timeout.connect(self.process_queue)
            self.timer.start(100)  # Check the queue every 100 ms

        def process_queue(self):
            while not self.queue.empty():
                url = self.queue.get()
                if url is None:  # Signal to close
                    self.close()
                    break
                self.add_image(url)

        def add_image(self, url):
            try:
                response = requests.get(url)
                response.raise_for_status()
                pixmap = QPixmap()
                pixmap.loadFromData(response.content)

                image_label = QLabel(self)
                image_label.setPixmap(pixmap)
                image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)

                # Scale image to fit, maintaining aspect ratio
                scaled_pixmap = pixmap.scaled(
                    200, 200, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation
                )
                image_label.setPixmap(scaled_pixmap)

                self.grid_layout.addWidget(image_label, self.row, self.col)

                # Update grid positions
                self.col += 1
                if self.col >= self.columns:
                    self.col = 0
                    self.row += 1

                # Auto-scroll to the bottom
                self.scroll_area.verticalScrollBar().setValue(
                    self.scroll_area.verticalScrollBar().maximum()
                )
            except Exception as e:
                print(f"Failed to load image from {url}: {e}")

    app = QApplication(sys.argv)
    viewer = ImageViewer(queue)
    viewer.show()
    sys.exit(app.exec())

class TrafficLogger:
    def __init__(self):
        # Set up a multiprocessing queue to share URLs
        self.queue = multiprocessing.Queue()

        # Start the GUI process
        self.gui_process = multiprocessing.Process(target=run_gui, args=(self.queue,))
        self.gui_process.start()

    def response(self, flow):
        content_type = flow.response.headers.get("Content-Type", "unknown")
        if content_type.startswith("image"):
            url = flow.request.url
            print(f'{content_type} {url}')
            self.queue.put(url)  # Send the URL to the GUI process

    def done(self):
        # Terminate the GUI process when mitmdump exits
        self.queue.put(None)  # Signal the GUI to close
        self.gui_process.join()

addons = [
    TrafficLogger()
]