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
 93        new_fields: list[tuple[KT, VT]] = []
 94        for field in self.fields:
 95            if self._kconv(field[0]) == key_kconv:
 96                if values:
 97                    new_fields.append((field[0], values.pop(0)))
 98            else:
 99                new_fields.append(field)
100        while values:
101            new_fields.append((key, values.pop(0)))
102        self.fields = tuple(new_fields)
103
104    def add(self, key: KT, value: VT) -> None:
105        """
106        Add an additional value for the given key at the bottom.
107        """
108        self.insert(len(self.fields), key, value)
109
110    def insert(self, index: int, key: KT, value: VT) -> None:
111        """
112        Insert an additional value for the given key at the specified position.
113        """
114        item = (key, value)
115        self.fields = self.fields[:index] + (item,) + self.fields[index:]
116
117    def keys(self, multi: bool = False):
118        """
119        Get all keys.
120
121        If `multi` is True, one key per value will be returned.
122        If `multi` is False, duplicate keys will only be returned once.
123        """
124        return (k for k, _ in self.items(multi))
125
126    def values(self, multi: bool = False):
127        """
128        Get all values.
129
130        If `multi` is True, all values will be returned.
131        If `multi` is False, only the first value per key will be returned.
132        """
133        return (v for _, v in self.items(multi))
134
135    def items(self, multi: bool = False):
136        """
137        Get all (key, value) tuples.
138
139        If `multi` is True, all `(key, value)` pairs will be returned.
140        If False, only one tuple per key is returned.
141        """
142        if multi:
143            return self.fields
144        else:
145            return super().items()
146
147
148class MultiDict(_MultiDict[KT, VT], serializable.Serializable):
149    """A concrete MultiDict, storing its own data."""
150
151    def __init__(self, fields=()):
152        super().__init__()
153        self.fields = tuple(tuple(i) for i in fields)  # type: ignore
154
155    @staticmethod
156    def _reduce_values(values):
157        return values[0]
158
159    @staticmethod
160    def _kconv(key):
161        return key
162
163    def get_state(self):
164        return self.fields
165
166    def set_state(self, state):
167        self.fields = tuple(tuple(x) for x in state)  # type: ignore
168
169    @classmethod
170    def from_state(cls, state):
171        return cls(state)
172
173
174class MultiDictView(_MultiDict[KT, VT]):
175    """
176    The MultiDictView provides the MultiDict interface over calculated data.
177    The view itself contains no state - data is retrieved from the parent on
178    request, and stored back to the parent on change.
179    """
180
181    def __init__(self, getter, setter):
182        self._getter = getter
183        self._setter = setter
184        super().__init__()
185
186    @staticmethod
187    def _kconv(key):
188        # All request-attributes are case-sensitive.
189        return key
190
191    @staticmethod
192    def _reduce_values(values):
193        # We just return the first element if
194        # multiple elements exist with the same key.
195        return values[0]
196
197    @property  # type: ignore
198    def fields(self):
199        return self._getter()
200
201    @fields.setter
202    def fields(self, value):
203        self._setter(value)
204
205    def copy(self) -> "MultiDict[KT,VT]":
206        return MultiDict(self.fields)