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

3 files changed, 11 insertions, 7 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"]