mitmproxy.types -> mitmproxy.coretypes

The types name is valuable, and we have a better use for it in collecting and
exposing types for options and commands.

The coretypes module should probably be split up anyway - it contains a
threading base class, a few container objects, and the defintion of our
serialization protocol. I was tempted to rename it to "uncagegorized" for the
sake of honesty.

5 files changed, 295 insertions, 0 deletions

diff --git a/mitmproxy/coretypes/__init__.py b/mitmproxy/coretypes/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/mitmproxy/coretypes/__init__.py
diff --git a/mitmproxy/coretypes/basethread.py b/mitmproxy/coretypes/basethread.py
new file mode 100644
index 00000000..a3c81d19
--- /dev/null
+++ b/mitmproxy/coretypes/basethread.py
@@ -0,0 +1,14 @@
+import time
+import threading
+
+
+class BaseThread(threading.Thread):
+    def __init__(self, name, *args, **kwargs):
+        super().__init__(name=name, *args, **kwargs)
+        self._thread_started = time.time()
+
+    def _threadinfo(self):
+        return "%s - age: %is" % (
+            self.name,
+            int(time.time() - self._thread_started)
+        )
diff --git a/mitmproxy/coretypes/bidi.py b/mitmproxy/coretypes/bidi.py
new file mode 100644
index 00000000..0982a34a
--- /dev/null
+++ b/mitmproxy/coretypes/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/mitmproxy/coretypes/multidict.py b/mitmproxy/coretypes/multidict.py
new file mode 100644
index 00000000..90f3013e
--- /dev/null
+++ b/mitmproxy/coretypes/multidict.py
@@ -0,0 +1,216 @@
+from abc import ABCMeta, abstractmethod
+
+from collections.abc import MutableMapping
+from mitmproxy.coretypes import serializable
+
+
+class _MultiDict(MutableMapping, metaclass=ABCMeta):
+    def __repr__(self):
+        fields = (
+            repr(field)
+            for field in self.fields
+        )
+        return "{cls}[{fields}]".format(
+            cls=type(self).__name__,
+            fields=", ".join(fields)
+        )
+
+    @staticmethod
+    @abstractmethod
+    def _reduce_values(values):
+        """
+        If a user accesses multidict["foo"], this method
+        reduces all values for "foo" to a single value that is returned.
+        For example, HTTP headers are folded, whereas we will just take
+        the first cookie we found with that name.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def _kconv(key):
+        """
+        This method converts a key to its canonical representation.
+        For example, HTTP headers are case-insensitive, so this method returns key.lower().
+        """
+
+    def __getitem__(self, key):
+        values = self.get_all(key)
+        if not values:
+            raise KeyError(key)
+        return self._reduce_values(values)
+
+    def __setitem__(self, key, value):
+        self.set_all(key, [value])
+
+    def __delitem__(self, key):
+        if key not in self:
+            raise KeyError(key)
+        key = self._kconv(key)
+        self.fields = tuple(
+            field for field in self.fields
+            if key != self._kconv(field[0])
+        )
+
+    def __iter__(self):
+        seen = set()
+        for key, _ in self.fields:
+            key_kconv = self._kconv(key)
+            if key_kconv not in seen:
+                seen.add(key_kconv)
+                yield key
+
+    def __len__(self):
+        return len(set(self._kconv(key) for key, _ in self.fields))
+
+    def __eq__(self, other):
+        if isinstance(other, MultiDict):
+            return self.fields == other.fields
+        return False
+
+    def get_all(self, key):
+        """
+        Return the list of all values for a given key.
+        If that key is not in the MultiDict, the return value will be an empty list.
+        """
+        key = self._kconv(key)
+        return [
+            value
+            for k, value in self.fields
+            if self._kconv(k) == key
+        ]
+
+    def set_all(self, key, values):
+        """
+        Remove the old values for a key and add new ones.
+        """
+        key_kconv = self._kconv(key)
+
+        new_fields = []
+        for field in self.fields:
+            if self._kconv(field[0]) == key_kconv:
+                if values:
+                    new_fields.append(
+                        (field[0], values.pop(0))
+                    )
+            else:
+                new_fields.append(field)
+        while values:
+            new_fields.append(
+                (key, values.pop(0))
+            )
+        self.fields = tuple(new_fields)
+
+    def add(self, key, value):
+        """
+        Add an additional value for the given key at the bottom.
+        """
+        self.insert(len(self.fields), key, value)
+
+    def insert(self, index, key, value):
+        """
+        Insert an additional value for the given key at the specified position.
+        """
+        item = (key, value)
+        self.fields = self.fields[:index] + (item,) + self.fields[index:]
+
+    def keys(self, multi=False):
+        """
+        Get all keys.
+
+        Args:
+            multi(bool):
+                If True, one key per value will be returned.
+                If False, duplicate keys will only be returned once.
+        """
+        return (
+            k
+            for k, _ in self.items(multi)
+        )
+
+    def values(self, multi=False):
+        """
+        Get all values.
+
+        Args:
+            multi(bool):
+                If True, all values will be returned.
+                If False, only the first value per key will be returned.
+        """
+        return (
+            v
+            for _, v in self.items(multi)
+        )
+
+    def items(self, multi=False):
+        """
+        Get all (key, value) tuples.
+
+        Args:
+            multi(bool):
+                If True, all (key, value) pairs will be returned
+                If False, only the first (key, value) pair per unique key will be returned.
+        """
+        if multi:
+            return self.fields
+        else:
+            return super().items()
+
+
+class MultiDict(_MultiDict, serializable.Serializable):
+    def __init__(self, fields=()):
+        super().__init__()
+        self.fields = tuple(
+            tuple(i) for i in fields
+        )
+
+    @staticmethod
+    def _reduce_values(values):
+        return values[0]
+
+    @staticmethod
+    def _kconv(key):
+        return key
+
+    def get_state(self):
+        return self.fields
+
+    def set_state(self, state):
+        self.fields = tuple(tuple(x) for x in state)
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(state)
+
+
+class MultiDictView(_MultiDict):
+    """
+    The MultiDictView provides the MultiDict interface over calculated data.
+    The view itself contains no state - data is retrieved from the parent on
+    request, and stored back to the parent on change.
+    """
+    def __init__(self, getter, setter):
+        self._getter = getter
+        self._setter = setter
+        super().__init__()
+
+    @staticmethod
+    def _kconv(key):
+        # All request-attributes are case-sensitive.
+        return key
+
+    @staticmethod
+    def _reduce_values(values):
+        # We just return the first element if
+        # multiple elements exist with the same key.
+        return values[0]
+
+    @property
+    def fields(self):
+        return self._getter()
+
+    @fields.setter
+    def fields(self, value):
+        self._setter(value)
+
+    def copy(self):
+        return MultiDict(self.fields)
diff --git a/mitmproxy/coretypes/serializable.py b/mitmproxy/coretypes/serializable.py
new file mode 100644
index 00000000..cd8539b0
--- /dev/null
+++ b/mitmproxy/coretypes/serializable.py
@@ -0,0 +1,36 @@
+import abc
+import uuid
+
+
+class Serializable(metaclass=abc.ABCMeta):
+    """
+    Abstract Base Class that defines an API to save an object's state and restore it later on.
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def from_state(cls, state):
+        """
+        Create a new object from the given state.
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_state(self):
+        """
+        Retrieve object state.
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_state(self, state):
+        """
+        Set object state to the given state.
+        """
+        raise NotImplementedError()
+
+    def copy(self):
+        state = self.get_state()
+        if isinstance(state, dict) and "id" in state:
+            state["id"] = str(uuid.uuid4())
+        return self.from_state(state)