Record and replay fixtures for your Python tests with mitmproxy

README.md

Record and replay fixtures in Python with mitmproxy

Why not use vcr?

I wrote this bercause vcr is currently not thread-safe, and it is less effort to mock everything if the record/replay infrastructure exists as an HTTP proxy rather than in-process.

Dockerfile

FROM python:3.8-buster
WORKDIR /code
# Don't run as root.
RUN groupadd -r baphomet \
    && useradd -m -d /home/baphomet -s /bin/bash -g baphomet baphomet

RUN wget -q https://snapshots.mitmproxy.org/5.0.1/mitmproxy-5.0.1-linux.tar.gz -O mitmproxy.tar.gz \
    && tar xvf mitmproxy.tar.gz \
    && mv mitmdump mitmproxy mitmweb /bin \
    && rm mitmproxy.tar.gz

# mitmproxy generates these certificates on its first run. I generated these
# on my host machine and am saving them in the container.
ADD mitmproxy_certificates /root/.mitmproxy
ADD mitmproxy_certificates /home/baphomet/.mitmproxy

# Add your other files here and install your Python packages.
# This depends on `certifi`.

# Allow scripts running as `baphomet` to edit the list of trusted
# TLS certificates, so you don't have to run your tests as root.
RUN chown baphomet:baphomet `python -c "import certifi; print(certifi.where())"`
USER baphomet

mitmcontext.py

import certifi
import os
import subprocess
import requests
import time


class MITMContext:
    """
    Context manager to record/replay unit tests with HTTP(S) requests.
    
    This assumes that you have `mitmdump` installed in your $PATH.
    This also assumes that this script has the permissions to edit
    the certificate store located at `certifi.where()`.

    Also, this wants port 8080 open by default, but that is easy to change.
    
    Usage:
        ```
        import requests
        # If some_file.mitm is not around, this will record the fixture
        # and write it to the file.
        with MITMContext(fixture_filename="some_file.mitm", replay=True):
            print("UUID is", requests.get("https://httpbin.org/uuid").text)
       # The fixture exists, so we don't talk to httpbin and instead read
       # from the file.
        with MITMContext(fixture_filename="some_file.mitm", replay=True):
            print("UUID is", requests.get("https://httpbin.org/uuid").text)
        ```
    """

    fixture_filename: str = ""
    replay: bool = False
    original_ca_root_contents: bytes = ""
    ca_cert_pem_filename: str = os.path.join(
        os.path.expanduser("~"), ".mitmproxy", "mitmproxy-ca-cert.pem"
    )
    ca_root: str = certifi.where()
    mitm_process = None
    replacements = [":~q:foo:bar"]

    def __init__(self, *, fixture_filename, replay=False, ca_cert_pem_filename=None):
        self.fixture_filename = fixture_filename
        self.replay = replay
        if ca_cert_pem_filename:
            self.ca_cert_pem_filename = ca_cert_pem_filename

    def __enter__(self):
        # Back up the original root certificate.
        with open(self.ca_root, "rb") as ca_root_handle:
            self.original_ca_root_contents = ca_root_handle.read()
        # Load up our MITM certificate.
        with open(self.ca_cert_pem_filename, "rb") as ca_cert_handle:
            ca_cert = ca_cert_handle.read()
        # Append our MITM certificate to the root≥
        with open(self.ca_root, "ab") as ca_root_handle:
            ca_root_handle.write(ca_cert)
        os.environ["HTTP_PROXY"] = "http://localhost:8080"
        os.environ["HTTPS_PROXY"] = "http://localhost:8080"
        if self.replay:
            self.mitmdump_replay()
        else:
            self.mitmdump_record()
        time.sleep(1)  # Wait for mitmproxy to accept connections
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        with open(self.ca_root, "wb") as ca_root_handle:
            ca_root_handle.write(self.original_ca_root_contents)
        self.mitm_process.terminate()

    def mitmdump_record(self):
        """Record a fixture to a filename."""
        self.mitm_process = subprocess.Popen(["mitmdump", "-w", self.fixture_filename])
        return self.mitm_process

    def mitmdump_replay(self):
        """Replay a fixture, but record if the fixture does not exist."""
        if not os.path.exists(self.fixture_filename):
            return self.mitmdump_record()
        self.mitm_process = subprocess.Popen(
            [
                "mitmdump",
                "-S",
                self.fixture_filename,
                "-k",
                "--set",
                "server_replay_nopop=true",
            ]
        )
        return self.mitm_process