Raiff1982 commited on
Commit
fc03166
·
verified ·
1 Parent(s): 89b7788

Update modules/secure_memory_loader.py

Browse files
Files changed (1) hide show
  1. modules/secure_memory_loader.py +33 -0
modules/secure_memory_loader.py CHANGED
@@ -2,6 +2,39 @@ import importlib.util
2
  from pathlib import Path
3
 
4
  def load_secure_memory_module(temp_path="/mnt/data/secure_memory.py"):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  """
6
  Dynamically loads secure_memory.py from a temporary path.
7
  Returns the module object.
 
2
  from pathlib import Path
3
 
4
  def load_secure_memory_module(temp_path="/mnt/data/secure_memory.py"):
5
+ SECURE_MEMORY_CODE = """
6
+ from cryptography.fernet import Fernet
7
+ import numpy as np
8
+ import base64
9
+
10
+ class SecureMemorySession:
11
+ def __init__(self, encryption_key: bytes = None):
12
+ self.key = encryption_key or Fernet.generate_key()
13
+ self.fernet = Fernet(self.key)
14
+ self.sessions = {}
15
+
16
+ def encrypt_vector(self, user_id: int, vector: np.ndarray) -> str:
17
+ vector_bytes = vector.tobytes()
18
+ encrypted = self.fernet.encrypt(vector_bytes)
19
+ encoded = base64.b64encode(encrypted).decode('utf-8')
20
+ if user_id not in self.sessions:
21
+ self.sessions[user_id] = []
22
+ self.sessions[user_id].append(encoded)
23
+ return encoded
24
+
25
+ def decrypt_vectors(self, user_id: int) -> list:
26
+ if user_id not in self.sessions:
27
+ return []
28
+ decrypted_vectors = []
29
+ for encoded in self.sessions[user_id]:
30
+ encrypted = base64.b64decode(encoded)
31
+ decrypted = self.fernet.decrypt(encrypted)
32
+ vector = np.frombuffer(decrypted, dtype=np.int64)
33
+ decrypted_vectors.append(vector)
34
+ return decrypted_vectors
35
+
36
+ def get_encryption_key(self) -> bytes:
37
+ return self.key
38
  """
39
  Dynamically loads secure_memory.py from a temporary path.
40
  Returns the module object.