netlib.utils.BiDi -> mitmproxy.types.bidi.BiDi

7 files changed, 53 insertions, 49 deletions

diff --git a/mitmproxy/types/bidi.py b/mitmproxy/types/bidi.py
new file mode 100644
index 00000000..0982a34a
--- /dev/null
+++ b/mitmproxy/types/bidi.py
@@ -0,0 +1,29 @@
+
+
+class BiDi:
+
+    """
+        A wee utility class for keeping bi-directional mappings, like field
+        constants in protocols. Names are attributes on the object, dict-like
+        access maps values to names:
+
+        CONST = BiDi(a=1, b=2)
+        assert CONST.a == 1
+        assert CONST.get_name(1) == "a"
+    """
+
+    def __init__(self, **kwargs):
+        self.names = kwargs
+        self.values = {}
+        for k, v in kwargs.items():
+            self.values[v] = k
+        if len(self.names) != len(self.values):
+            raise ValueError("Duplicate values not allowed.")
+
+    def __getattr__(self, k):
+        if k in self.names:
+            return self.names[k]
+        raise AttributeError("No such attribute: %s", k)
+
+    def get_name(self, n, default=None):
+        return self.values.get(n, default)
diff --git a/netlib/socks.py b/netlib/socks.py
index 9f1adb98..30fe1c9d 100644
--- a/netlib/socks.py
+++ b/netlib/socks.py
@@ -3,6 +3,7 @@ import array
 import ipaddress
 
 from netlib import tcp, utils
+from mitmproxy.types import bidi
 
 
 class SocksError(Exception):
@@ -10,24 +11,24 @@ class SocksError(Exception):
         super().__init__(message)
         self.code = code
 
-VERSION = utils.BiDi(
+VERSION = bidi.BiDi(
     SOCKS4=0x04,
     SOCKS5=0x05
 )
 
-CMD = utils.BiDi(
+CMD = bidi.BiDi(
     CONNECT=0x01,
     BIND=0x02,
     UDP_ASSOCIATE=0x03
 )
 
-ATYP = utils.BiDi(
+ATYP = bidi.BiDi(
     IPV4_ADDRESS=0x01,
     DOMAINNAME=0x03,
     IPV6_ADDRESS=0x04
 )
 
-REP = utils.BiDi(
+REP = bidi.BiDi(
     SUCCEEDED=0x00,
     GENERAL_SOCKS_SERVER_FAILURE=0x01,
     CONNECTION_NOT_ALLOWED_BY_RULESET=0x02,
@@ -39,14 +40,14 @@ REP = utils.BiDi(
     ADDRESS_TYPE_NOT_SUPPORTED=0x08,
 )
 
-METHOD = utils.BiDi(
+METHOD = bidi.BiDi(
     NO_AUTHENTICATION_REQUIRED=0x00,
     GSSAPI=0x01,
     USERNAME_PASSWORD=0x02,
     NO_ACCEPTABLE_METHODS=0xFF
 )
 
-USERNAME_PASSWORD_VERSION = utils.BiDi(
+USERNAME_PASSWORD_VERSION = bidi.BiDi(
     DEFAULT=0x01
 )
 
diff --git a/netlib/utils.py b/netlib/utils.py
index 12b94d74..779eaa27 100644
--- a/netlib/utils.py
+++ b/netlib/utils.py
@@ -16,35 +16,6 @@ def getbit(byte, offset):
     return bool(byte & mask)
 
 
-class BiDi:
-
-    """
-        A wee utility class for keeping bi-directional mappings, like field
-        constants in protocols. Names are attributes on the object, dict-like
-        access maps values to names:
-
-        CONST = BiDi(a=1, b=2)
-        assert CONST.a == 1
-        assert CONST.get_name(1) == "a"
-    """
-
-    def __init__(self, **kwargs):
-        self.names = kwargs
-        self.values = {}
-        for k, v in kwargs.items():
-            self.values[v] = k
-        if len(self.names) != len(self.values):
-            raise ValueError("Duplicate values not allowed.")
-
-    def __getattr__(self, k):
-        if k in self.names:
-            return self.names[k]
-        raise AttributeError("No such attribute: %s", k)
-
-    def get_name(self, n, default=None):
-        return self.values.get(n, default)
-
-
 _label_valid = re.compile(b"(?!-)[A-Z\d-]{1,63}(?<!-)$", re.IGNORECASE)
 
 
diff --git a/netlib/websockets/frame.py b/netlib/websockets/frame.py
index e022a95c..26303094 100644
--- a/netlib/websockets/frame.py
+++ b/netlib/websockets/frame.py
@@ -6,6 +6,7 @@ from netlib import tcp
 from mitmproxy.utils import strutils
 from netlib import utils
 from mitmproxy.utils import human
+from mitmproxy.types import bidi
 from .masker import Masker
 
 
@@ -15,7 +16,7 @@ MAX_64_BIT_INT = (1 << 64)
 DEFAULT = object()
 
 # RFC 6455, Section 5.2 - Base Framing Protocol
-OPCODE = utils.BiDi(
+OPCODE = bidi.BiDi(
     CONTINUE=0x00,
     TEXT=0x01,
     BINARY=0x02,
@@ -25,7 +26,7 @@ OPCODE = utils.BiDi(
 )
 
 # RFC 6455, Section 7.4.1 - Defined Status Codes
-CLOSE_REASON = utils.BiDi(
+CLOSE_REASON = bidi.BiDi(
     NORMAL_CLOSURE=1000,
     GOING_AWAY=1001,
     PROTOCOL_ERROR=1002,
diff --git a/pathod/protocols/http2.py b/pathod/protocols/http2.py
index 4dcc4b76..838469d6 100644
--- a/pathod/protocols/http2.py
+++ b/pathod/protocols/http2.py
@@ -4,11 +4,11 @@ import time
 import hyperframe.frame
 from hpack.hpack import Encoder, Decoder
 
-from netlib import utils
 from netlib.http import http2
 import netlib.http.headers
 import netlib.http.response
 import netlib.http.request
+from mitmproxy.types import bidi
 
 from .. import language
 
@@ -22,7 +22,7 @@ class TCPHandler:
 
 class HTTP2StateProtocol:
 
-    ERROR_CODES = utils.BiDi(
+    ERROR_CODES = bidi.BiDi(
         NO_ERROR=0x0,
         PROTOCOL_ERROR=0x1,
         INTERNAL_ERROR=0x2,
diff --git a/test/mitmproxy/test_types_bidi.py b/test/mitmproxy/test_types_bidi.py
new file mode 100644
index 00000000..a012c568
--- /dev/null
+++ b/test/mitmproxy/test_types_bidi.py
@@ -0,0 +1,11 @@
+from mitmproxy.types import bidi
+from netlib import tutils
+
+
+def test_bidi():
+    b = bidi.BiDi(a=1, b=2)
+    assert b.a == 1
+    assert b.get_name(1) == "a"
+    assert b.get_name(5) is None
+    tutils.raises(AttributeError, getattr, b, "c")
+    tutils.raises(ValueError, bidi.BiDi, one=1, two=1)
diff --git a/test/netlib/test_utils.py b/test/netlib/test_utils.py
index 9dcbffd8..db33fa31 100644
--- a/test/netlib/test_utils.py
+++ b/test/netlib/test_utils.py
@@ -1,6 +1,6 @@
 # coding=utf-8
 
-from netlib import utils, tutils
+from netlib import utils
 
 
 def test_is_valid_host():
@@ -8,12 +8,3 @@ def test_is_valid_host():
     assert utils.is_valid_host(b"one.two")
     assert not utils.is_valid_host(b"one" * 255)
     assert utils.is_valid_host(b"one.two.")
-
-
-def test_bidi():
-    b = utils.BiDi(a=1, b=2)
-    assert b.a == 1
-    assert b.get_name(1) == "a"
-    assert b.get_name(5) is None
-    tutils.raises(AttributeError, getattr, b, "c")
-    tutils.raises(ValueError, utils.BiDi, one=1, two=1)