mitmproxy.coretypes.multidict

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