mitmproxy.coretypes.multidict

  1from abc import ABCMeta
  2from abc import abstractmethod
  3from collections.abc import Iterator, MutableMapping, Sequence
  4from typing import TypeVar
  5
  6from mitmproxy.coretypes import serializable
  7
  8KT = TypeVar("KT")
  9VT = TypeVar("VT")
 10
 11
 12class _MultiDict(MutableMapping[KT, VT], metaclass=ABCMeta):
 13    """
 14    A MultiDict is a dictionary-like data structure that supports multiple values per key.
 15    """
 16
 17    fields: tuple[tuple[KT, VT], ...]
 18    """The underlying raw datastructure."""
 19
 20    def __repr__(self):
 21        fields = (repr(field) for field in self.fields)
 22        return "{cls}[{fields}]".format(
 23            cls=type(self).__name__, fields=", ".join(fields)
 24        )
 25
 26    @staticmethod
 27    @abstractmethod
 28    def _reduce_values(values: Sequence[VT]) -> VT:
 29        """
 30        If a user accesses multidict["foo"], this method
 31        reduces all values for "foo" to a single value that is returned.
 32        For example, HTTP headers are folded, whereas we will just take
 33        the first cookie we found with that name.
 34        """
 35
 36    @staticmethod
 37    @abstractmethod
 38    def _kconv(key: KT) -> KT:
 39        """
 40        This method converts a key to its canonical representation.
 41        For example, HTTP headers are case-insensitive, so this method returns key.lower().
 42        """
 43
 44    def __getitem__(self, key: KT) -> VT:
 45        values = self.get_all(key)
 46        if not values:
 47            raise KeyError(key)
 48        return self._reduce_values(values)
 49
 50    def __setitem__(self, key: KT, value: VT) -> None:
 51        self.set_all(key, [value])
 52
 53    def __delitem__(self, key: KT) -> None:
 54        if key not in self:
 55            raise KeyError(key)
 56        key = self._kconv(key)
 57        self.fields = tuple(
 58            field for field in self.fields if key != self._kconv(field[0])
 59        )
 60
 61    def __iter__(self) -> Iterator[KT]:
 62        seen = set()
 63        for key, _ in self.fields:
 64            key_kconv = self._kconv(key)
 65            if key_kconv not in seen:
 66                seen.add(key_kconv)
 67                yield key
 68
 69    def __len__(self) -> int:
 70        return len({self._kconv(key) for key, _ in self.fields})
 71
 72    def __eq__(self, other) -> bool:
 73        if isinstance(other, MultiDict):
 74            return self.fields == other.fields
 75        return False
 76
 77    def get_all(self, key: KT) -> list[VT]:
 78        """
 79        Return the list of all values for a given key.
 80        If that key is not in the MultiDict, the return value will be an empty list.
 81        """
 82        key = self._kconv(key)
 83        return [value for k, value in self.fields if self._kconv(k) == key]
 84
 85    def set_all(self, key: KT, values: list[VT]) -> None:
 86        """
 87        Remove the old values for a key and add new ones.
 88        """
 89        key_kconv = self._kconv(key)
 90
 91        new_fields: list[tuple[KT, VT]] = []
 92        for field in self.fields:
 93            if self._kconv(field[0]) == key_kconv:
 94                if values:
 95                    new_fields.append((field[0], values.pop(0)))
 96            else:
 97                new_fields.append(field)
 98        while values:
 99            new_fields.append((key, values.pop(0)))
100        self.fields = tuple(new_fields)
101
102    def add(self, key: KT, value: VT) -> None:
103        """
104        Add an additional value for the given key at the bottom.
105        """
106        self.insert(len(self.fields), key, value)
107
108    def insert(self, index: int, key: KT, value: VT) -> None:
109        """
110        Insert an additional value for the given key at the specified position.
111        """
112        item = (key, value)
113        self.fields = self.fields[:index] + (item,) + self.fields[index:]
114
115    def keys(self, multi: bool = False):
116        """
117        Get all keys.
118
119        If `multi` is True, one key per value will be returned.
120        If `multi` is False, duplicate keys will only be returned once.
121        """
122        return (k for k, _ in self.items(multi))
123
124    def values(self, multi: bool = False):
125        """
126        Get all values.
127
128        If `multi` is True, all values will be returned.
129        If `multi` is False, only the first value per key will be returned.
130        """
131        return (v for _, v in self.items(multi))
132
133    def items(self, multi: bool = False):
134        """
135        Get all (key, value) tuples.
136
137        If `multi` is True, all `(key, value)` pairs will be returned.
138        If False, only one tuple per key is returned.
139        """
140        if multi:
141            return self.fields
142        else:
143            return super().items()
144
145
146class MultiDict(_MultiDict[KT, VT], serializable.Serializable):
147    """A concrete MultiDict, storing its own data."""
148
149    def __init__(self, fields=()):
150        super().__init__()
151        self.fields = tuple(tuple(i) for i in fields)
152
153    @staticmethod
154    def _reduce_values(values):
155        return values[0]
156
157    @staticmethod
158    def _kconv(key):
159        return key
160
161    def get_state(self):
162        return self.fields
163
164    def set_state(self, state):
165        self.fields = tuple(tuple(x) for x in state)
166
167    @classmethod
168    def from_state(cls, state):
169        return cls(state)
170
171
172class MultiDictView(_MultiDict[KT, VT]):
173    """
174    The MultiDictView provides the MultiDict interface over calculated data.
175    The view itself contains no state - data is retrieved from the parent on
176    request, and stored back to the parent on change.
177    """
178
179    def __init__(self, getter, setter):
180        self._getter = getter
181        self._setter = setter
182        super().__init__()
183
184    @staticmethod
185    def _kconv(key):
186        # All request-attributes are case-sensitive.
187        return key
188
189    @staticmethod
190    def _reduce_values(values):
191        # We just return the first element if
192        # multiple elements exist with the same key.
193        return values[0]
194
195    @property  # type: ignore
196    def fields(self):
197        return self._getter()
198
199    @fields.setter
200    def fields(self, value):
201        self._setter(value)
202
203    def copy(self) -> "MultiDict[KT,VT]":
204        return MultiDict(self.fields)