# Copyright (C) 2022 DigeeX
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Data structures used in Raider.
"""
from typing import Any, Dict, Iterator, List, Optional, Tuple
from raider.plugins.basic import Cookie, Header
[docs]class DataStore:
"""Class defining a dictionary-like data structure.
This class was created to hold information relevant to Raider in a
structure similar to Python dictionaries.
"""
def __init__(self, data: Optional[Dict[Any, Any]]) -> None:
"""Initializes the DataStore object.
Given a dictionary with the data, store them in this object.
Args:
data:
A dictionary with Any elements to be stored.
"""
self._index = -1
if data:
self._store = data
else:
self._store = {}
[docs] def __getitem__(self, key: Any) -> Any:
"""Getter to return an element with the key."""
if key in self._store:
return self._store[key]
return None
[docs] def __setitem__(self, key: Any, value: Any) -> None:
"""Setter to add a new element to DataStore."""
self._store.update({key: value})
[docs] def __iter__(self) -> Iterator[Any]:
"""Iterator to yield the keys."""
for key in list(self._store):
yield key
[docs] def __next__(self) -> Any:
"""Iterator to get the next element."""
self._index += 1
if self._index >= len(self._store):
self._index = -1
raise StopIteration
key = list(self._store)[self._index]
return self._store[key]
[docs] def update(self, data: Dict[Any, Any]) -> None:
"""Updates the DataStore with a new element."""
self._store.update(data)
[docs] def pop(self, name: Any) -> Any:
"""Pops an element from the DataStore."""
return self._store.pop(name)
[docs] def list_keys(self) -> List[Any]:
"""Returns a list of the keys in the DataStore."""
return list(self._store)
[docs] def list_values(self) -> List[Any]:
"""Returns a list of the values in the DataStore."""
data = []
for key in self._store:
data.append(self._store[key])
return data
[docs] def to_dict(self) -> Dict[Any, Any]:
"""Returns the DataStore elements as a dictionary."""
return self._store
[docs] def items(self) -> List[Tuple[Any, Any]]:
"""Returns a list of tuples containing the keys and values."""
data = []
for key in self._store:
data.append((key, self._store[key]))
return data
[docs]class CookieStore(DataStore):
"""Class storing the HTTP cookies.
This class inherits from DataStore, and converts the values into
Cookie objects.
"""
def __init__(self, data: Optional[List[Cookie]]) -> None:
"""Initializes a CookieStore object.
Given a list of Cookie objects, create the CookieStore
containing them.
Args:
data:
A list of Cookies to be added to the CookieStore.
"""
values = {}
if data:
for cookie in data:
values[cookie.name] = cookie
super().__init__(values)
[docs] def set(self, cookie: Cookie) -> None:
"""Sets the value of a Cookie.
Given a Cookie object, add or update its value in the
CookieStore.
Args:
cookie:
A Cookie object to be added to the CookieStore
"""
super().update({cookie.name: cookie.value})
[docs] def merge(self, cookiestore: "CookieStore") -> None:
"""Merge CookieStore object with another one."""
for item in cookiestore:
self._store[item] = cookiestore[item]
[docs] @classmethod
def from_dict(cls, data: Optional[Dict[str, str]]) -> "CookieStore":
"""Creates a CookieStore object from a dictionary.
Given a dictionary with cookie values, creates a CookieStore
object and returns it.
Args:
data:
A dictionary with cookie values. Those will be mapped in
Cookie objects.
Returns:
A CookieStore object containing the cookies created from the
supplied dictionary.
"""
cookielist = []
if data:
for name, value in data.items():
cookie = Cookie(name, value)
cookielist.append(cookie)
return cls(cookielist)