Spaces:
Runtime error
Runtime error
File size: 1,192 Bytes
ac8274a 5c9ac41 ac8274a 5c9ac41 ac8274a 5c9ac41 ac8274a 5c9ac41 ac8274a 5c9ac41 |
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 |
import base64
import json
from cryptography.fernet import Fernet
class SecureMemorySession:
"""Handles encryption/decryption for user vectors"""
def __init__(self, key: bytes):
self.key = base64.urlsafe_b64encode(key[:32].ljust(32, b'0')) # Fernet key must be 32 bytes
self.fernet = Fernet(self.key)
self.memory_store = {}
def encrypt_vector(self, user_id, vector):
"""Encrypts and stores vector under user ID"""
if user_id not in self.memory_store:
self.memory_store[user_id] = []
serialized = json.dumps(vector.tolist()).encode()
encrypted = self.fernet.encrypt(serialized)
self.memory_store[user_id].append(encrypted)
def decrypt_vectors(self, user_id):
"""Decrypts and returns all stored vectors for a user"""
if user_id not in self.memory_store:
return []
decrypted = []
for enc in self.memory_store[user_id]:
try:
raw = self.fernet.decrypt(enc)
array = json.loads(raw.decode())
decrypted.append(array)
except Exception:
continue
return decrypted |