diff options
| -rw-r--r-- | mitmproxy/contrib/tnetstring.py | 14 | ||||
| -rw-r--r-- | mitmproxy/flow/io.py | 7 | ||||
| -rw-r--r-- | mitmproxy/flow/io_compat.py | 76 | ||||
| -rw-r--r-- | mitmproxy/models/connections.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_contrib_tnetstring.py | 8 |
5 files changed, 66 insertions, 41 deletions
diff --git a/mitmproxy/contrib/tnetstring.py b/mitmproxy/contrib/tnetstring.py index 5fc26b45..0383f98e 100644 --- a/mitmproxy/contrib/tnetstring.py +++ b/mitmproxy/contrib/tnetstring.py @@ -96,7 +96,7 @@ def _rdumpq(q, size, value): elif value is False: write(b'5:false!') return size + 8 - elif isinstance(value, int): + elif isinstance(value, six.integer_types): data = str(value).encode() ldata = len(data) span = str(ldata).encode() @@ -119,7 +119,7 @@ def _rdumpq(q, size, value): write(b'%s:%s,' % (span, data)) return size + 2 + len(span) + ldata elif isinstance(value, six.text_type): - data = value.encode() + data = value.encode("utf8") ldata = len(data) span = str(ldata).encode() write(b'%s:%s;' % (span, data)) @@ -137,8 +137,6 @@ def _rdumpq(q, size, value): write(b'}') init_size = size = size + 1 for (k, v) in value.items(): - if isinstance(k, str): - k = k.encode("ascii", "strict") size = _rdumpq(q, size, v) size = _rdumpq(q, size, k) span = str(size - init_size).encode() @@ -184,13 +182,17 @@ def load(file_handle): def parse(data_type, data): + if six.PY2: + data_type = ord(data_type) # type: (int, bytes) -> TSerializable if data_type == ord(b','): return data if data_type == ord(b';'): - return data.decode() + return data.decode("utf8") if data_type == ord(b'#'): try: + if six.PY2: + return long(data) return int(data) except ValueError: raise ValueError("not a tnetstring: invalid integer literal: {}".format(data)) @@ -220,8 +222,6 @@ def parse(data_type, data): d = {} while data: key, data = pop(data) - if isinstance(key, bytes): - key = key.decode("ascii", "strict") val, data = pop(data) d[key] = val return d diff --git a/mitmproxy/flow/io.py b/mitmproxy/flow/io.py index e5716940..276d7a5b 100644 --- a/mitmproxy/flow/io.py +++ b/mitmproxy/flow/io.py @@ -44,10 +44,9 @@ class FlowReader: raise exceptions.FlowReadException(str(e)) if can_tell: off = self.fo.tell() - data_type = data["type"].decode() - if data_type not in models.FLOW_TYPES: - raise exceptions.FlowReadException("Unknown flow type: {}".format(data_type)) - yield models.FLOW_TYPES[data_type].from_state(data) + if data["type"] not in models.FLOW_TYPES: + raise exceptions.FlowReadException("Unknown flow type: {}".format(data["type"])) + yield models.FLOW_TYPES[data["type"]].from_state(data) except ValueError: # Error is due to EOF if can_tell and self.fo.tell() == off and self.fo.read() == b'': diff --git a/mitmproxy/flow/io_compat.py b/mitmproxy/flow/io_compat.py index 55971f5e..1e67dde4 100644 --- a/mitmproxy/flow/io_compat.py +++ b/mitmproxy/flow/io_compat.py @@ -3,44 +3,74 @@ This module handles the import of mitmproxy flows generated by old versions. """ from __future__ import absolute_import, print_function, division +import six + from netlib import version def convert_013_014(data): - data["request"]["first_line_format"] = data["request"].pop("form_in") - data["request"]["http_version"] = "HTTP/" + ".".join(str(x) for x in data["request"].pop("httpversion")) - data["response"]["http_version"] = "HTTP/" + ".".join(str(x) for x in data["response"].pop("httpversion")) - data["response"]["status_code"] = data["response"].pop("code") - data["response"]["body"] = data["response"].pop("content") - data["server_conn"].pop("state") - data["server_conn"]["via"] = None - data["version"] = (0, 14) + data[b"request"][b"first_line_format"] = data[b"request"].pop(b"form_in") + data[b"request"][b"http_version"] = b"HTTP/" + ".".join(str(x) for x in data[b"request"].pop(b"httpversion")).encode() + data[b"response"][b"http_version"] = b"HTTP/" + ".".join(str(x) for x in data[b"response"].pop(b"httpversion")).encode() + data[b"response"][b"status_code"] = data[b"response"].pop(b"code") + data[b"response"][b"body"] = data[b"response"].pop(b"content") + data[b"server_conn"].pop(b"state") + data[b"server_conn"][b"via"] = None + data[b"version"] = (0, 14) return data def convert_014_015(data): - data["version"] = (0, 15) + data[b"version"] = (0, 15) return data def convert_015_016(data): - for m in ("request", "response"): - if "body" in data[m]: - data[m]["content"] = data[m].pop("body") - if "msg" in data["response"]: - data["response"]["reason"] = data["response"].pop("msg") - data["request"].pop("form_out", None) - data["version"] = (0, 16) + for m in (b"request", b"response"): + if b"body" in data[m]: + data[m][b"content"] = data[m].pop(b"body") + if b"msg" in data[b"response"]: + data[b"response"][b"reason"] = data[b"response"].pop(b"msg") + data[b"request"].pop(b"form_out", None) + data[b"version"] = (0, 16) return data def convert_016_017(data): - data["server_conn"]["peer_address"] = None - data["version"] = (0, 17) + data[b"server_conn"][b"peer_address"] = None + data[b"version"] = (0, 17) return data def convert_017_018(data): + if not six.PY2: + # Python 2 versions of mitmproxy did not support serializing unicode. + def dict_keys_to_str(o): + if isinstance(o, dict): + return {k.decode(): dict_keys_to_str(v) for k, v in o.items()} + else: + return o + data = dict_keys_to_str(data) + + def dict_vals_to_str(o, decode): + for k, v in decode.items(): + if not o or k not in o: + continue + if v is True: + o[k] = o[k].decode() + else: + dict_vals_to_str(o[k], v) + dict_vals_to_str(data, { + "type": True, + "id": True, + "request": { + "first_line_format": True + }, + "error": { + "msg": True + } + }) + data["server_conn"]["ip_address"] = data["server_conn"].pop("peer_address") data["version"] = (0, 18) return data @@ -57,13 +87,13 @@ converters = { def migrate_flow(flow_data): while True: - flow_version = tuple(flow_data["version"][:2]) - if flow_version == version.IVERSION[:2]: + flow_version = tuple(flow_data.get(b"version", flow_data.get("version"))) + if flow_version[:2] == version.IVERSION[:2]: break - elif flow_version in converters: - flow_data = converters[flow_version](flow_data) + elif flow_version[:2] in converters: + flow_data = converters[flow_version[:2]](flow_data) else: - v = ".".join(str(i) for i in flow_data["version"]) + v = ".".join(str(i) for i in flow_version) raise ValueError( "{} cannot read files serialized with version {}.".format(version.MITMPROXY, v) ) diff --git a/mitmproxy/models/connections.py b/mitmproxy/models/connections.py index d71379bc..3e1a0928 100644 --- a/mitmproxy/models/connections.py +++ b/mitmproxy/models/connections.py @@ -162,7 +162,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject): source_address=tcp.Address, ssl_established=bool, cert=certutils.SSLCert, - sni=str, + sni=bytes, timestamp_start=float, timestamp_tcp_setup=float, timestamp_ssl_setup=float, diff --git a/test/mitmproxy/test_contrib_tnetstring.py b/test/mitmproxy/test_contrib_tnetstring.py index 8ae35a25..908cec27 100644 --- a/test/mitmproxy/test_contrib_tnetstring.py +++ b/test/mitmproxy/test_contrib_tnetstring.py @@ -15,7 +15,9 @@ FORMAT_EXAMPLES = { {'hello': [12345678901, b'this', True, None, b'\x00\x00\x00\x00']}, b'5:12345#': 12345, b'12:this is cool,': b'this is cool', + b'19:this is unicode \xe2\x98\x85;': u'this is unicode \u2605', b'0:,': b'', + b'0:;': u'', b'0:~': None, b'4:true!': True, b'5:false!': False, @@ -78,12 +80,6 @@ class Test_Format(unittest.TestCase): self.assertEqual(v, tnetstring.loads(tnetstring.dumps(v))) self.assertEqual((v, b""), tnetstring.pop(tnetstring.dumps(v))) - def test_unicode_handling(self): - with self.assertRaises(ValueError): - tnetstring.dumps(u"hello") - self.assertEqual(tnetstring.dumps(u"hello".encode()), b"5:hello,") - self.assertEqual(type(tnetstring.loads(b"5:hello,")), bytes) - def test_roundtrip_format_unicode(self): for _ in range(500): v = get_random_object() |