23581258074

Committed 26 Mar 2026 06:45AM UTC coverage: 86.264% (-7.7%) from 93.96%

Build # 23581258074

Build Type

Pull #41

github

Committed by

web-flow

Commit Message

Merge 78cf0ad5c into a189234b3

Pull Request Pull Request #41: Modernize project: CLI, dependency removal, restructure

Coverage Stats

401 of 457 new or added lines in 12 files covered. (87.75%)

22 existing lines in 1 file now uncovered.

628 of 728 relevant lines covered (86.26%)

4.31 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

82.76

/mitm/proxy.py

"""
Man-in-the-middle.
"""

from __future__ import annotations

import asyncio
import logging
import pathlib
from typing import List, Optional

import OpenSSL

from mitm import __data__
from mitm.extension.middleware import Log
from mitm.extension.protocol import HTTP, InvalidProtocol
from mitm.models import Connection, Host, Middleware, Protocol
from mitm.utils.crypto import CertificateAuthority

TEMPLATES = pathlib.Path(__file__).parent / "templates"
CERT_PAGE = (TEMPLATES / "cert.html").read_bytes()

logger = logging.getLogger(__package__)
logging.getLogger("asyncio").setLevel(logging.CRITICAL)


class MITM:
    """
    Man-in-the-middle proxy server.

    Example:

        .. code-block:: python

            from mitm import MITM

            mitm = MITM()
            mitm.run()

        .. code-block:: python

            async with MITM() as mitm:
                ...
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 8888,
        protocols: Optional[List[Protocol]] = None,
        middlewares: Optional[List[Middleware]] = None,
        certificate_authority: Optional[CertificateAuthority] = None,
    ):
        """
        Initializes the MITM class.

        Args:
            host: Host to listen on. Defaults to `127.0.0.1`.
            port: Port to listen on. Defaults to `8888`.
            protocols: List of protocols to use. Defaults to `[protocol.HTTP]`.
            middlewares: List of middlewares to use. Defaults to `[middleware.Log]`.
            certificate_authority: Certificate authority to use. Defaults to `CertificateAuthority()`.
        """
        self.host = host
        self.port = port
        self.certificate_authority = certificate_authority if certificate_authority else CertificateAuthority()

        # Stores the CA certificate and private key.
        cert_path, key_path = __data__ / "mitm.crt", __data__ / "mitm.key"
        self.certificate_authority.save(cert_path=cert_path, key_path=key_path)

        # Initialize any middleware that is not already initialized.
        middlewares = middlewares if middlewares else [Log]
        new_middlewares = []
        for middleware in middlewares:
            if isinstance(middleware, type):
                middleware = middleware()
            new_middlewares.append(middleware)
        self.middlewares = new_middlewares

        # Initialize any protocol that is not already initialized.
        protocols = protocols if protocols else [HTTP]
        new_protocols = []
        for protocol in protocols:
            if isinstance(protocol, type):
                protocol = protocol(
                    certificate_authority=self.certificate_authority,
                    middlewares=self.middlewares,
                )
            new_protocols.append(protocol)
        self.protocols = new_protocols

        self.server: asyncio.Server | None = None

    async def start(self):
        """
        Start the MITM proxy server.

        Notes:
            Begins accepting connections immediately. Use `run` for blocking usage,
            or the class as an async context manager for automatic cleanup.

        Raises:
            OSError: If the server cannot bind to the host and port.
        """
        self.server = await asyncio.start_server(
            lambda reader, writer: self.mitm(
                Connection(
                    client=Host(reader=reader, writer=writer),
                    server=Host(),
                )
            ),
            host=self.host,
            port=self.port,
        )

        for middleware in self.middlewares:
            await middleware.mitm_started(host=self.host, port=self.port)

    async def stop(self):
        """
        Stop the MITM proxy server.
        """
        if self.server:
            self.server.close()
            await self.server.wait_closed()
            self.server = None

    def run(self):
        """
        Run the MITM proxy server (blocking).

        Notes:
            Starts the server and serves forever until interrupted.
            For non-blocking usage, use `start` and `stop` directly.
        """

        async def serve():
            await self.start()
            async with self.server:
                await self.server.serve_forever()

        asyncio.run(serve())

    async def __aenter__(self) -> MITM:
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
        await self.stop()
        return False

    async def serve_direct(self, connection: Connection, data: bytes):
        """
        Serve the cert download page for direct (non-proxy) requests.
        """
        first_line = data.split(b"\r\n", 1)[0]
        parts = first_line.split(b" ")
        path = parts[1].decode() if len(parts) > 1 else "/"

        if path == "/":
            body = CERT_PAGE
            content_type = "text/html; charset=utf-8"
            disposition = ""
        elif path == "/cert.pem":
            body = OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_PEM, self.certificate_authority.cert)
            content_type = "application/x-pem-file"
            disposition = "Content-Disposition: attachment; filename=mitm-ca.pem\r\n"
        elif path in ("/cert.cer", "/cert.crt"):
            body = OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_ASN1, self.certificate_authority.cert)
            ext = path.split(".")[-1]
            content_type = "application/x-x509-ca-cert"
            disposition = f"Content-Disposition: attachment; filename=mitm-ca.{ext}\r\n"
        else:
            body = b"404 Not Found"
            content_type = "text/plain"
            disposition = ""

        header = (
            f"HTTP/1.1 200 OK\r\n"
            f"Content-Type: {content_type}\r\n"
            f"Content-Length: {len(body)}\r\n"
            f"{disposition}"
            f"Connection: close\r\n"
            f"\r\n"
        )
        connection.client.writer.write(header.encode() + body)
        await connection.client.writer.drain()
        connection.client.writer.close()
        await connection.client.writer.wait_closed()

    async def mitm(self, connection: Connection):
        """
        Handles an incoming connection (single connection).

        Warning:
            This method is not intended to be called directly.
        """

        #  Calls middlewares for client initial connect.
        for middleware in self.middlewares:
            await middleware.client_connected(connection=connection)

        # Gets the bytes needed to identify the protocol.
        min_bytes_needed = max(proto.bytes_needed for proto in self.protocols)
        data = await connection.client.reader.read(n=min_bytes_needed)

        # Calls middleware on client's data.
        for middleware in self.middlewares:
            data = await middleware.client_data(connection=connection, data=data)

        # Direct requests (GET /path, not GET http://...) are served by the cert page.
        first_line = data.split(b"\r\n", 1)[0]
        if first_line.startswith(b"GET /") and not first_line.startswith(b"GET http"):
            await self.serve_direct(connection, data)
            return

        # Finds the protocol that matches the data.
        proto = None
        for prtcl in self.protocols:
            proto = prtcl
            try:
                # Attempts to resolve the protocol, and connect to the server.
                host, port, tls = await proto.resolve(connection=connection, data=data)
                await proto.connect(connection=connection, host=host, port=port, tls=tls, data=data)
            except InvalidProtocol:  # pragma: no cover
                proto = None
            else:
                # Stop searching for working protocols.
                break

        # Protocol was found, and we connected to a server.
        if proto and connection.server:
            # Sets the connection protocol.
            connection.protocol = proto

            # Calls middleware for server initial connect.
            for middleware in self.middlewares:
                await middleware.server_connected(connection=connection)

            # Handles the data between the client and server.
            await proto.handle(connection=connection)

        # Protocol identified, but we did not connect to a server.
        elif proto and not connection.server:  # pragma: no cover
            raise ValueError(
                "The protocol was found, but the server was not connected to succesfully. "
                f"Check the {proto.__class__.__name__} implementation."
            )

        # No protocol was found for the data.
        else:  # pragma: no cover
            raise ValueError("No protocol was found. Check the protocols list.")

        # If a server connection exists after handling it, we close it.
        if connection.server and connection.server.mitm_managed:
            connection.server.writer.close()
            await connection.server.writer.wait_closed()

            # Calls the server's 'disconnected' middleware.
            for middleware in self.middlewares:
                await middleware.server_disconnected(connection=connection)

        # Attempts to disconnect with the client.
        # In some instances 'wait_closed()' might hang. This is a known issue that
        # happens when and if the client keeps the connection alive, and, unfortunately,
        # there is nothing we can do about it. This is a reported bug in asyncio.
        # https://bugs.python.org/issue39758
        if connection.client and connection.client.mitm_managed:
            connection.client.writer.close()
            await connection.client.writer.wait_closed()

            # Calls the client 'disconnected' middleware.
            for middleware in self.middlewares:
                await middleware.client_disconnected(connection=connection)

1	"""
2	Man-in-the-middle.
3	"""
4
5	from __future__ import annotations	5✔
6
7	import asyncio	5✔
8	import logging	5✔
9	import pathlib	5✔
10	from typing import List, Optional	5✔
11
12	import OpenSSL	5✔
13
14	from mitm import __data__	5✔
15	from mitm.extension.middleware import Log	5✔
16	from mitm.extension.protocol import HTTP, InvalidProtocol	5✔
17	from mitm.models import Connection, Host, Middleware, Protocol	5✔
18	from mitm.utils.crypto import CertificateAuthority	5✔
19
20	TEMPLATES = pathlib.Path(__file__).parent / "templates"	5✔
21	CERT_PAGE = (TEMPLATES / "cert.html").read_bytes()	5✔
22
23	logger = logging.getLogger(__package__)	5✔
24	logging.getLogger("asyncio").setLevel(logging.CRITICAL)	5✔
25
26
27	class MITM:	5✔
28	"""
29	Man-in-the-middle proxy server.
30
31	Example:
32
33	.. code-block:: python
34
35	from mitm import MITM
36
37	mitm = MITM()
38	mitm.run()
39
40	.. code-block:: python
41
42	async with MITM() as mitm:
43	...
44	"""
45
46	def __init__(	5✔
47	self,
48	host: str = "127.0.0.1",
49	port: int = 8888,
50	protocols: Optional[List[Protocol]] = None,
51	middlewares: Optional[List[Middleware]] = None,
52	certificate_authority: Optional[CertificateAuthority] = None,
53	):
54	"""
55	Initializes the MITM class.
56
57	Args:
58	host: Host to listen on. Defaults to `127.0.0.1`.
59	port: Port to listen on. Defaults to `8888`.
60	protocols: List of protocols to use. Defaults to `[protocol.HTTP]`.
61	middlewares: List of middlewares to use. Defaults to `[middleware.Log]`.
62	certificate_authority: Certificate authority to use. Defaults to `CertificateAuthority()`.
63	"""
64	self.host = host	5✔
65	self.port = port	5✔
66	self.certificate_authority = certificate_authority if certificate_authority else CertificateAuthority()	5✔
67
68	# Stores the CA certificate and private key.
69	cert_path, key_path = __data__ / "mitm.crt", __data__ / "mitm.key"	5✔
70	self.certificate_authority.save(cert_path=cert_path, key_path=key_path)	5✔
71
72	# Initialize any middleware that is not already initialized.
73	middlewares = middlewares if middlewares else [Log]	5✔
74	new_middlewares = []	5✔
75	for middleware in middlewares:	5✔
76	if isinstance(middleware, type):	5✔
77	middleware = middleware()	5✔
78	new_middlewares.append(middleware)	5✔
79	self.middlewares = new_middlewares	5✔
80
81	# Initialize any protocol that is not already initialized.
82	protocols = protocols if protocols else [HTTP]	5✔
83	new_protocols = []	5✔
84	for protocol in protocols:	5✔
85	if isinstance(protocol, type):	5✔
86	protocol = protocol(	5✔
87	certificate_authority=self.certificate_authority,
88	middlewares=self.middlewares,
89	)
90	new_protocols.append(protocol)	5✔
91	self.protocols = new_protocols	5✔
92
93	self.server: asyncio.Server \| None = None	5✔
94
95	async def start(self):	5✔
96	"""
97	Start the MITM proxy server.
98
99	Notes:
100	Begins accepting connections immediately. Use `run` for blocking usage,
101	or the class as an async context manager for automatic cleanup.
102
103	Raises:
104	OSError: If the server cannot bind to the host and port.
105	"""
106	self.server = await asyncio.start_server(	5✔
107	lambda reader, writer: self.mitm(
108	Connection(
109	client=Host(reader=reader, writer=writer),
110	server=Host(),
111	)
112	),
113	host=self.host,
114	port=self.port,
115	)
116
117	for middleware in self.middlewares:	5✔
118	await middleware.mitm_started(host=self.host, port=self.port)	5✔
119
120	async def stop(self):	5✔
121	"""
122	Stop the MITM proxy server.
123	"""
124	if self.server:	5✔
125	self.server.close()	5✔
126	await self.server.wait_closed()	5✔
127	self.server = None	5✔
128
129	def run(self):	5✔
130	"""
131	Run the MITM proxy server (blocking).
132
133	Notes:
134	Starts the server and serves forever until interrupted.
135	For non-blocking usage, use `start` and `stop` directly.
136	"""
137
NEW 138	async def serve():	×
NEW 139	await self.start()	×
NEW 140	async with self.server:	×
NEW 141	await self.server.serve_forever()	×
142
NEW 143	asyncio.run(serve())	×
144
145	async def __aenter__(self) -> MITM:	5✔
146	await self.start()	5✔
147	return self	5✔
148
149	async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:	5✔
150	await self.stop()	5✔
151	return False	5✔
152
153	async def serve_direct(self, connection: Connection, data: bytes):	5✔
154	"""
155	Serve the cert download page for direct (non-proxy) requests.
156	"""
157	first_line = data.split(b"\r\n", 1)[0]	5✔
158	parts = first_line.split(b" ")	5✔
159	path = parts[1].decode() if len(parts) > 1 else "/"	5✔
160
161	if path == "/":	5✔
162	body = CERT_PAGE	5✔
163	content_type = "text/html; charset=utf-8"	5✔
164	disposition = ""	5✔
165	elif path == "/cert.pem":	5✔
166	body = OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_PEM, self.certificate_authority.cert)	5✔
167	content_type = "application/x-pem-file"	5✔
168	disposition = "Content-Disposition: attachment; filename=mitm-ca.pem\r\n"	5✔
169	elif path in ("/cert.cer", "/cert.crt"):	5✔
170	body = OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_ASN1, self.certificate_authority.cert)	5✔
171	ext = path.split(".")[-1]	5✔
172	content_type = "application/x-x509-ca-cert"	5✔
173	disposition = f"Content-Disposition: attachment; filename=mitm-ca.{ext}\r\n"	5✔
174	else:
175	body = b"404 Not Found"	5✔
176	content_type = "text/plain"	5✔
177	disposition = ""	5✔
178
179	header = (	5✔
180	f"HTTP/1.1 200 OK\r\n"
181	f"Content-Type: {content_type}\r\n"
182	f"Content-Length: {len(body)}\r\n"
183	f"{disposition}"
184	f"Connection: close\r\n"
185	f"\r\n"
186	)
187	connection.client.writer.write(header.encode() + body)	5✔
188	await connection.client.writer.drain()	5✔
189	connection.client.writer.close()	5✔
190	await connection.client.writer.wait_closed()	5✔
191
192	async def mitm(self, connection: Connection):	5✔
193	"""
194	Handles an incoming connection (single connection).
195
196	Warning:
197	This method is not intended to be called directly.
198	"""
199
200	# Calls middlewares for client initial connect.
201	for middleware in self.middlewares:	5✔
202	await middleware.client_connected(connection=connection)	5✔
203
204	# Gets the bytes needed to identify the protocol.
205	min_bytes_needed = max(proto.bytes_needed for proto in self.protocols)	5✔
206	data = await connection.client.reader.read(n=min_bytes_needed)	5✔
207
208	# Calls middleware on client's data.
209	for middleware in self.middlewares:	5✔
210	data = await middleware.client_data(connection=connection, data=data)	5✔
211
212	# Direct requests (GET /path, not GET http://...) are served by the cert page.
213	first_line = data.split(b"\r\n", 1)[0]	5✔
214	if first_line.startswith(b"GET /") and not first_line.startswith(b"GET http"):	5✔
215	await self.serve_direct(connection, data)	5✔
216	return	5✔
217
218	# Finds the protocol that matches the data.
219	proto = None	5✔
220	for prtcl in self.protocols:	5✔
221	proto = prtcl	5✔
222	try:	5✔
223	# Attempts to resolve the protocol, and connect to the server.
224	host, port, tls = await proto.resolve(connection=connection, data=data)	5✔
225	await proto.connect(connection=connection, host=host, port=port, tls=tls, data=data)	5✔
226	except InvalidProtocol: # pragma: no cover
227	proto = None
228	else:
229	# Stop searching for working protocols.
230	break	×
231
232	# Protocol was found, and we connected to a server.
233	if proto and connection.server:	5✔
234	# Sets the connection protocol.
235	connection.protocol = proto	×
236
237	# Calls middleware for server initial connect.
238	for middleware in self.middlewares:	×
239	await middleware.server_connected(connection=connection)	×
240
241	# Handles the data between the client and server.
242	await proto.handle(connection=connection)	×
243
244	# Protocol identified, but we did not connect to a server.
245	elif proto and not connection.server: # pragma: no cover
246	raise ValueError(
247	"The protocol was found, but the server was not connected to succesfully. "
248	f"Check the {proto.__class__.__name__} implementation."
249	)
250
251	# No protocol was found for the data.
252	else: # pragma: no cover
253	raise ValueError("No protocol was found. Check the protocols list.")
254
255	# If a server connection exists after handling it, we close it.
256	if connection.server and connection.server.mitm_managed:	×
257	connection.server.writer.close()	×
258	await connection.server.writer.wait_closed()	×
259
260	# Calls the server's 'disconnected' middleware.
261	for middleware in self.middlewares:	×
262	await middleware.server_disconnected(connection=connection)	×
263
264	# Attempts to disconnect with the client.
265	# In some instances 'wait_closed()' might hang. This is a known issue that
266	# happens when and if the client keeps the connection alive, and, unfortunately,
267	# there is nothing we can do about it. This is a reported bug in asyncio.
268	# https://bugs.python.org/issue39758
269	if connection.client and connection.client.mitm_managed:	×
270	connection.client.writer.close()	×
271	await connection.client.writer.wait_closed()	×
272
273	# Calls the client 'disconnected' middleware.
274	for middleware in self.middlewares:	×
275	await middleware.client_disconnected(connection=connection)	×

synchronizing / mitm / 23581258074

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous