Spaces:
Runtime error
Runtime error
File size: 3,712 Bytes
2609fac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import os
import base64
import streamlit as st
from typing import Tuple
from typing import Optional
from cryptography import fernet
from typing import MutableMapping
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from utilities_cookies.cookie_manager import CookieManager
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
@st.cache_data
def key_from_parameters(salt: bytes, iterations: int, password: str):
kdf = PBKDF2HMAC(
algorithm=hashes.SHA256(),
length=32,
salt=salt,
iterations=iterations,
)
return base64.urlsafe_b64encode(kdf.derive(password.encode('utf-8')))
class EncryptedCookieManager(MutableMapping[str, str]):
def __init__(
self, *,
password: str,
path: str = None,
prefix: str = "",
key_params_cookie="EncryptedCookieManager.key_params",
ignore_broken=True,
):
self._cookie_manager = CookieManager(path=path, prefix=prefix)
self._fernet: Optional[Fernet] = None
self._key_params_cookie = key_params_cookie
self._password = password
self._ignore_broken = ignore_broken
def ready(self):
return self._cookie_manager.ready()
def save(self):
return self._cookie_manager.save()
def _encrypt(self, value):
self._setup_fernet()
return self._fernet.encrypt(value)
def _decrypt(self, value):
self._setup_fernet()
return self._fernet.decrypt(value)
def _setup_fernet(self):
if self._fernet is not None:
return
key_params = self._get_key_params()
if not key_params:
key_params = self._initialize_new_key_params()
salt, iterations, magic = key_params
key = key_from_parameters(
salt=salt,
iterations=iterations,
password=self._password
)
self._fernet = Fernet(key)
def _get_key_params(self) -> Optional[Tuple[bytes, int, bytes]]:
raw_key_params = self._cookie_manager.get(self._key_params_cookie)
if not raw_key_params:
return
try:
raw_salt, raw_iterations, raw_magic = raw_key_params.split(':')
return base64.b64decode(raw_salt), int(raw_iterations), base64.b64decode(raw_magic)
except (ValueError, TypeError):
print(f"Failed to parse key parameters from cookie {raw_key_params}")
return
def _initialize_new_key_params(self) -> Tuple[bytes, int, bytes]:
salt = os.urandom(16)
iterations = 390000
magic = os.urandom(16)
self._cookie_manager[self._key_params_cookie] = b':'.join([
base64.b64encode(salt),
str(iterations).encode('ascii'),
base64.b64encode(magic)
]).decode('ascii')
return salt, iterations, magic
def __repr__(self):
if self.ready():
return f'<EncryptedCookieManager: {dict(self)!r}>'
return '<EncryptedCookieManager: not ready>'
def __getitem__(self, k: str) -> str or None:
try:
return self._decrypt(self._cookie_manager[k].encode('utf-8')).decode('utf-8')
except fernet.InvalidToken:
if self._ignore_broken:
return
raise
def __iter__(self):
return iter(self._cookie_manager)
def __len__(self):
return len(self._cookie_manager)
def __setitem__(self, key: str, value: str) -> None:
self._cookie_manager[key] = self._encrypt(value.encode('utf-8')).decode('utf-8')
def __delitem__(self, key: str) -> None:
del self._cookie_manager[key]
|