diff options
| -rw-r--r-- | mitmproxy/connections.py | 23 | ||||
| -rw-r--r-- | mitmproxy/flow.py | 1 | ||||
| -rw-r--r-- | mitmproxy/io_compat.py | 22 | ||||
| -rw-r--r-- | mitmproxy/test/tflow.py | 3 | ||||
| -rw-r--r-- | mitmproxy/types/serializable.py | 6 | ||||
| -rw-r--r-- | mitmproxy/version.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_connections.py | 25 | ||||
| -rw-r--r-- | test/mitmproxy/types/test_serializable.py | 19 |
8 files changed, 93 insertions, 8 deletions
diff --git a/mitmproxy/connections.py b/mitmproxy/connections.py index 9359b67d..01721a71 100644 --- a/mitmproxy/connections.py +++ b/mitmproxy/connections.py @@ -1,6 +1,7 @@ import time import os +import uuid from mitmproxy import stateobject from mitmproxy import certs @@ -41,6 +42,7 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject): self.clientcert = None self.ssl_established = None + self.id = str(uuid.uuid4()) self.mitmcert = None self.timestamp_start = time.time() self.timestamp_end = None @@ -73,6 +75,14 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject): port=self.address[1], ) + def __eq__(self, other): + if isinstance(other, ClientConnection): + return self.id == other.id + return False + + def __hash__(self): + return hash(self.id) + @property def tls_established(self): return self.ssl_established @@ -82,6 +92,7 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject): self.ssl_established = value _stateobject_attributes = dict( + id=str, address=tuple, ssl_established=bool, clientcert=certs.SSLCert, @@ -110,6 +121,7 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject): @classmethod def make_dummy(cls, address): return cls.from_state(dict( + id=str(uuid.uuid4()), address=address, clientcert=None, mitmcert=None, @@ -165,6 +177,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject): def __init__(self, address, source_address=None, spoof_source_address=None): tcp.TCPClient.__init__(self, address, source_address, spoof_source_address) + self.id = str(uuid.uuid4()) self.alpn_proto_negotiated = None self.tls_version = None self.via = None @@ -196,6 +209,14 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject): port=self.address[1], ) + def __eq__(self, other): + if isinstance(other, ServerConnection): + return self.id == other.id + return False + + def __hash__(self): + return hash(self.id) + @property def tls_established(self): return self.ssl_established @@ -205,6 +226,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject): self.ssl_established = value _stateobject_attributes = dict( + id=str, address=tuple, ip_address=tuple, source_address=tuple, @@ -228,6 +250,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject): @classmethod def make_dummy(cls, address): return cls.from_state(dict( + id=str(uuid.uuid4()), address=address, ip_address=address, cert=None, diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 73d6090e..7ad71516 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -112,7 +112,6 @@ class Flow(stateobject.StateObject): def copy(self): f = super().copy() - f.id = str(uuid.uuid4()) f.live = False return f diff --git a/mitmproxy/io_compat.py b/mitmproxy/io_compat.py index 8fa7b77c..7d839ffd 100644 --- a/mitmproxy/io_compat.py +++ b/mitmproxy/io_compat.py @@ -1,7 +1,7 @@ """ This module handles the import of mitmproxy flows generated by old versions. """ - +import uuid from typing import Any, Dict from mitmproxy import version @@ -113,6 +113,25 @@ def convert_300_4(data): return data +client_connections = {} +server_connections = {} + + +def convert_4_5(data): + data["version"] = 5 + client_conn_key = ( + data["client_conn"]["timestamp_start"], + *data["client_conn"]["address"] + ) + server_conn_key = ( + data["server_conn"]["timestamp_start"], + *data["server_conn"]["source_address"] + ) + data["client_conn"]["id"] = client_connections.setdefault(client_conn_key, str(uuid.uuid4())) + data["server_conn"]["id"] = server_connections.setdefault(server_conn_key, str(uuid.uuid4())) + return data + + def _convert_dict_keys(o: Any) -> Any: if isinstance(o, dict): return {strutils.always_str(k): _convert_dict_keys(v) for k, v in o.items()} @@ -164,6 +183,7 @@ converters = { (1, 0): convert_100_200, (2, 0): convert_200_300, (3, 0): convert_300_4, + 4: convert_4_5, } diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py index 7fbe1727..270021cb 100644 --- a/mitmproxy/test/tflow.py +++ b/mitmproxy/test/tflow.py @@ -1,4 +1,5 @@ import io +import uuid from mitmproxy.net import websockets from mitmproxy.test import tutils @@ -146,6 +147,7 @@ def tclient_conn(): @return: mitmproxy.proxy.connection.ClientConnection """ c = connections.ClientConnection.from_state(dict( + id=str(uuid.uuid4()), address=("address", 22), clientcert=None, mitmcert=None, @@ -169,6 +171,7 @@ def tserver_conn(): @return: mitmproxy.proxy.connection.ServerConnection """ c = connections.ServerConnection.from_state(dict( + id=str(uuid.uuid4()), address=("address", 22), source_address=("address", 22), ip_address=None, diff --git a/mitmproxy/types/serializable.py b/mitmproxy/types/serializable.py index 49892ffc..cd8539b0 100644 --- a/mitmproxy/types/serializable.py +++ b/mitmproxy/types/serializable.py @@ -1,4 +1,5 @@ import abc +import uuid class Serializable(metaclass=abc.ABCMeta): @@ -29,4 +30,7 @@ class Serializable(metaclass=abc.ABCMeta): raise NotImplementedError() def copy(self): - return self.from_state(self.get_state()) + state = self.get_state() + if isinstance(state, dict) and "id" in state: + state["id"] = str(uuid.uuid4()) + return self.from_state(state) diff --git a/mitmproxy/version.py b/mitmproxy/version.py index 2882d1fb..006ec868 100644 --- a/mitmproxy/version.py +++ b/mitmproxy/version.py @@ -5,7 +5,7 @@ MITMPROXY = "mitmproxy " + VERSION # Serialization format version. This is displayed nowhere, it just needs to be incremented by one # for each change the the file format. -FLOW_FORMAT_VERSION = 4 +FLOW_FORMAT_VERSION = 5 if __name__ == "__main__": print(VERSION) diff --git a/test/mitmproxy/test_connections.py b/test/mitmproxy/test_connections.py index 0083f57c..67a6552f 100644 --- a/test/mitmproxy/test_connections.py +++ b/test/mitmproxy/test_connections.py @@ -66,8 +66,18 @@ class TestClientConnection: assert c.timestamp_start == 42 c3 = c.copy() + assert c3.get_state() != c.get_state() + c.id = c3.id = "foo" assert c3.get_state() == c.get_state() + def test_eq(self): + c = tflow.tclient_conn() + c2 = c.copy() + assert c == c + assert c != c2 + assert c != 42 + assert hash(c) != hash(c2) + class TestServerConnection: @@ -147,6 +157,21 @@ class TestServerConnection: with pytest.raises(ValueError, matches='sni must be str, not '): c.establish_ssl(None, b'foobar') + def test_state(self): + c = tflow.tserver_conn() + c2 = c.copy() + assert c2.get_state() != c.get_state() + c.id = c2.id = "foo" + assert c2.get_state() == c.get_state() + + def test_eq(self): + c = tflow.tserver_conn() + c2 = c.copy() + assert c == c + assert c != c2 + assert c != 42 + assert hash(c) != hash(c2) + class TestClientConnectionTLS: diff --git a/test/mitmproxy/types/test_serializable.py b/test/mitmproxy/types/test_serializable.py index dd4a3778..390d17e1 100644 --- a/test/mitmproxy/types/test_serializable.py +++ b/test/mitmproxy/types/test_serializable.py @@ -1,3 +1,5 @@ +import copy + from mitmproxy.types import serializable @@ -6,17 +8,17 @@ class SerializableDummy(serializable.Serializable): self.i = i def get_state(self): - return self.i + return copy.copy(self.i) def set_state(self, i): self.i = i - def from_state(self, state): - return type(self)(state) + @classmethod + def from_state(cls, state): + return cls(state) class TestSerializable: - def test_copy(self): a = SerializableDummy(42) assert a.i == 42 @@ -26,3 +28,12 @@ class TestSerializable: a.set_state(1) assert a.i == 1 assert b.i == 42 + + def test_copy_id(self): + a = SerializableDummy({ + "id": "foo", + "foo": 42 + }) + b = a.copy() + assert a.get_state()["id"] != b.get_state()["id"] + assert a.get_state()["foo"] == b.get_state()["foo"] |