tnetstring: keys are str on py3. migrate flow.io_compat

7 files changed, 20 insertions, 16 deletions

diff --git a/mitmproxy/contrib/py3/tnetstring.py b/mitmproxy/contrib/py3/tnetstring.py
index 6f38a245..6998fc82 100644
--- a/mitmproxy/contrib/py3/tnetstring.py
+++ b/mitmproxy/contrib/py3/tnetstring.py
@@ -126,6 +126,8 @@ def _rdumpq(q: collections.deque, size: int, value: TSerializable) -> int:
         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()
@@ -154,7 +156,7 @@ def load(file_handle: io.BinaryIO) -> TSerializable:
     #  Note that the netstring spec explicitly forbids padding zeros.
     c = file_handle.read(1)
     data_length = b""
-    while ord(b'0') <= ord(c) <= ord(b'9'):
+    while c.isdigit():
         data_length += c
         if len(data_length) > 9:
             raise ValueError("not a tnetstring: absurdly large length prefix")
@@ -202,6 +204,8 @@ def parse(data_type: int, data: bytes) -> TSerializable:
         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
@@ -230,4 +234,4 @@ def pop(data: bytes) -> Tuple[TSerializable, bytes]:
     return parse(data_type, data), remain
 
 
-__all__ = ["dump", "dumps", "load", "loads"]
+__all__ = ["dump", "dumps", "load", "loads", "pop"]
diff --git a/mitmproxy/contrib/py3/tnetstring_tests.py b/mitmproxy/contrib/py3/tnetstring_tests.py
index 545889c8..4ee184d5 100644
--- a/mitmproxy/contrib/py3/tnetstring_tests.py
+++ b/mitmproxy/contrib/py3/tnetstring_tests.py
@@ -11,7 +11,7 @@ FORMAT_EXAMPLES = {
     b'0:}': {},
     b'0:]': [],
     b'51:5:hello,39:11:12345678901#4:this,4:true!0:~4:\x00\x00\x00\x00,]}':
-            {b'hello': [12345678901, b'this', True, None, b'\x00\x00\x00\x00']},
+            {'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'0:,': b'',
@@ -41,7 +41,7 @@ def get_random_object(random=random, depth=0):
             d = {}
             for _ in range(n):
                 n = random.randint(0,100)
-                k = bytes([random.randint(32,126) for _ in range(n)])
+                k = str([random.randint(32,126) for _ in range(n)])
                 d[k] = get_random_object(random,depth+1)
             return d
     else:
diff --git a/mitmproxy/contrib/tnetstring.py b/mitmproxy/contrib/tnetstring.py
index 58daec5c..1ebaba21 100644
--- a/mitmproxy/contrib/tnetstring.py
+++ b/mitmproxy/contrib/tnetstring.py
@@ -1,8 +1,8 @@
 import six
 
 if six.PY2:
-    from .py2.tnetstring import load, loads, dump, dumps
+    from .py2.tnetstring import load, loads, dump, dumps, pop
 else:
-    from .py3.tnetstring import load, loads, dump, dumps
+    from .py3.tnetstring import load, loads, dump, dumps, pop
 
-__all__ = ["load", "loads", "dump", "dumps"]
+__all__ = ["load", "loads", "dump", "dumps", "pop"]
diff --git a/mitmproxy/flow/io.py b/mitmproxy/flow/io.py
index 671ddf43..e5716940 100644
--- a/mitmproxy/flow/io.py
+++ b/mitmproxy/flow/io.py
@@ -44,12 +44,13 @@ class FlowReader:
                     raise exceptions.FlowReadException(str(e))
                 if can_tell:
                     off = self.fo.tell()
-                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)
+                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)
         except ValueError:
             # Error is due to EOF
-            if can_tell and self.fo.tell() == off and self.fo.read() == '':
+            if can_tell and self.fo.tell() == off and self.fo.read() == b'':
                 return
             raise exceptions.FlowReadException("Invalid data format.")
 
diff --git a/mitmproxy/flow/io_compat.py b/mitmproxy/flow/io_compat.py
index 1023e87f..55971f5e 100644
--- a/mitmproxy/flow/io_compat.py
+++ b/mitmproxy/flow/io_compat.py
@@ -9,6 +9,7 @@ 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")
@@ -26,8 +27,6 @@ def convert_015_016(data):
     for m in ("request", "response"):
         if "body" in data[m]:
             data[m]["content"] = data[m].pop("body")
-        if "httpversion" in data[m]:
-            data[m]["http_version"] = data[m].pop("httpversion")
     if "msg" in data["response"]:
         data["response"]["reason"] = data["response"].pop("msg")
     data["request"].pop("form_out", None)
diff --git a/test/mitmproxy/test_contrib_tnetstring.py b/test/mitmproxy/test_contrib_tnetstring.py
index 17654ad9..8ae35a25 100644
--- a/test/mitmproxy/test_contrib_tnetstring.py
+++ b/test/mitmproxy/test_contrib_tnetstring.py
@@ -12,7 +12,7 @@ FORMAT_EXAMPLES = {
     b'0:}': {},
     b'0:]': [],
     b'51:5:hello,39:11:12345678901#4:this,4:true!0:~4:\x00\x00\x00\x00,]}':
-    {b'hello': [12345678901, b'this', True, None, b'\x00\x00\x00\x00']},
+    {'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'0:,': b'',
@@ -43,7 +43,7 @@ def get_random_object(random=random, depth=0):
             d = {}
             for _ in range(n):
                 n = random.randint(0, 100)
-                k = bytes([random.randint(32, 126) for _ in range(n)])
+                k = str([random.randint(32, 126) for _ in range(n)])
                 d[k] = get_random_object(random, depth + 1)
             return d
     else:
diff --git a/tox.ini b/tox.ini
index a7b5e7d3..251609a5 100644
--- a/tox.ini
+++ b/tox.ini
@@ -16,7 +16,7 @@ commands =
 
 [testenv:py35]
 setenv =
-  TESTS = test/netlib test/pathod/ test/mitmproxy/script test/mitmproxy/test_contentview.py test/mitmproxy/test_custom_contentview.py test/mitmproxy/test_app.py test/mitmproxy/test_controller.py test/mitmproxy/test_fuzzing.py test/mitmproxy/test_script.py test/mitmproxy/test_web_app.py test/mitmproxy/test_utils.py test/mitmproxy/test_stateobject.py test/mitmproxy/test_cmdline.py test/mitmproxy/test_contrib_tnetstring.py test/mitmproxy/test_proxy.py test/mitmproxy/test_protocol_http1.py test/mitmproxy/test_platform_pf.py test/mitmproxy/test_server.py test/mitmproxy/test_filt.py test/mitmproxy/test_flow_export.py test/mitmproxy/test_web_master.py
+  TESTS = test/netlib test/pathod/ test/mitmproxy/script test/mitmproxy/test_contentview.py test/mitmproxy/test_custom_contentview.py test/mitmproxy/test_app.py test/mitmproxy/test_controller.py test/mitmproxy/test_fuzzing.py test/mitmproxy/test_script.py test/mitmproxy/test_web_app.py test/mitmproxy/test_utils.py test/mitmproxy/test_stateobject.py test/mitmproxy/test_cmdline.py test/mitmproxy/test_contrib_tnetstring.py test/mitmproxy/test_proxy.py test/mitmproxy/test_protocol_http1.py test/mitmproxy/test_platform_pf.py test/mitmproxy/test_server.py test/mitmproxy/test_filt.py test/mitmproxy/test_flow_export.py test/mitmproxy/test_web_master.py test/mitmproxy/test_flow_format_compat.py
   HOME = {envtmpdir}
 
 [testenv:docs]