allow multiple private keys

Author: mahbd

Framework/Built_In_Automation/Shared_Resources/secrets.py

@@ -16,8 +16,7 @@
 import json
 import base64
 import inspect
-from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Dict, Optional
 
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.asymmetric import padding
@@ -28,6 +27,9 @@
 from Framework.Built_In_Automation.Shared_Resources import BuiltInFunctionSharedResources as sr
 
 
+from settings import ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR
+
+
 MODULE_NAME = inspect.getmodulename(__file__)
 
 
@@ -44,7 +46,7 @@ class Secret:
 
     def __init__(self):
         self._cache: Dict[str, str] = {}
-        self._private_key_path = Path.home() / "zeuz_node_downloads" / "private_key.pem"
+        self._private_key_folder = ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR
 
     def __getitem__(self, key_name: str) -> str:
         """
@@ -148,24 +150,32 @@ def __getitem__(self, key_name: str) -> str:
             raise KeyError(f"Failed to retrieve secret '{key_name}': {str(e)}")
 
 
-    def _load_private_key(self):
+    def _load_private_keys(self):
         sModuleInfo = inspect.currentframe().f_code.co_name + " : " + MODULE_NAME
 
         try:
-            if not self._private_key_path.exists():
+            if not self._private_key_folder.exists():
                 CommonUtil.ExecLog(
                     sModuleInfo,
-                    f"Private key not found at {self._private_key_path}",
+                    f"Private key folder not found at {self._private_key_folder}",
                     3
                 )
-                raise FileNotFoundError(f"Private key not found at {self._private_key_path}")
+                raise FileNotFoundError(f"Private key folder not found at {self._private_key_folder}")
+            
+            private_keys = []
+            for pem_file in self._private_key_folder.glob("*.pem"):
+                with open(pem_file, 'rb') as f:
+                    private_key = serialization.load_pem_private_key(f.read(), password=None)
+                    private_keys.append(private_key)
+            
+            if not private_keys:
+                raise FileNotFoundError(f"No .pem files found in {self._private_key_folder}")
 
-            with open(self._private_key_path, 'rb') as f:
-                return serialization.load_pem_private_key(f.read(), password=None)
+            return private_keys
         except Exception as e:
             CommonUtil.ExecLog(
                 sModuleInfo,
-                f"Failed to load private key: {str(e)}",
+                f"Failed to load private keys: {str(e)}",
                 3
             )
             raise
@@ -180,32 +190,38 @@ def _decrypt_data(self, encrypted_data: str) -> str:
         Returns:
             Decrypted plaintext string
         """
-        private_key = self._load_private_key()
+        private_keys = self._load_private_keys()
 
         decoded_data = base64.b64decode(encrypted_data)
         data = json.loads(decoded_data.decode('utf-8'))
 
         encrypted_aes_key = base64.b64decode(data['encryptedKey'])
-        aes_key = private_key.decrypt(
-            encrypted_aes_key,
-            padding.OAEP(
-                mgf=padding.MGF1(algorithm=hashes.SHA256()),
-                algorithm=hashes.SHA256(),
-                label=None
-            )
-        )
-        
         iv = base64.b64decode(data['iv'])
         encrypted_content = base64.b64decode(data['encryptedData'])
 
-        cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
-        decryptor = cipher.decryptor()
-        decrypted_padded = decryptor.update(encrypted_content) + decryptor.finalize()
-        
-        unpadder = PKCS7(128).unpadder()
-        decrypted = unpadder.update(decrypted_padded) + unpadder.finalize()
+        for private_key in private_keys:
+            try:
+                aes_key = private_key.decrypt(
+                    encrypted_aes_key,
+                    padding.OAEP(
+                        mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                        algorithm=hashes.SHA256(),
+                        label=None
+                    )
+                )
+                
+                cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
+                decryptor = cipher.decryptor()
+                decrypted_padded = decryptor.update(encrypted_content) + decryptor.finalize()
+                
+                unpadder = PKCS7(128).unpadder()
+                decrypted = unpadder.update(decrypted_padded) + unpadder.finalize()
+                
+                return decrypted.decode('utf-8')
+            except Exception:
+                continue
 
-        return decrypted.decode('utf-8')
+        raise Exception("No private key could decrypt the data")
 
     def clear_cache(self, key_name: Optional[str] = None):
         """

Framework/Built_In_Automation/Web/Selenium/utils.py

@@ -18,10 +18,7 @@
 import struct
 import urllib.request
 from rich.progress import Progress
-
-
-# ZeuZ Node Downloads base directory
-ZEUZ_NODE_DOWNLOADS_DIR = Path.home() / "zeuz_node_downloads"
+from settings import ZEUZ_NODE_DOWNLOADS_DIR
 
 
 class ChromeForTesting:

settings.py

@@ -7,6 +7,10 @@
 # BASE_DIR or PROJECT_ROOT or Zeuz_Python_Node dir
 PROJECT_ROOT = Path(__file__).parent
 
+ZEUZ_NODE_ARTIFACTS_DIR = Path.home() / "zeuz_node_artifacts"
+ZEUZ_NODE_DOWNLOADS_DIR = ZEUZ_NODE_ARTIFACTS_DIR / "zeuz_node_downloads"
+ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR = ZEUZ_NODE_ARTIFACTS_DIR / "rsa_private_keys"
+
 # AutomationLog dir
 AutomationLog_DIR = PROJECT_ROOT / "AutomationLog"