diff --git a/skyvern/forge/sdk/encrypt/aes.py b/skyvern/forge/sdk/encrypt/aes.py index 31784a96d4..2c3574374a 100644 --- a/skyvern/forge/sdk/encrypt/aes.py +++ b/skyvern/forge/sdk/encrypt/aes.py @@ -17,17 +17,24 @@ def __init__(self, *, secret_key: str, salt: str | None = None, iv: str | None = self.salt = hashlib.md5(salt.encode("utf-8")).digest() if salt else default_salt self.iv = hashlib.md5(iv.encode("utf-8")).digest() if iv else default_iv + + # Precompute derived key for this instance, since secret_key and salt do not change + self._derived_key: bytes | None = None + def method(self) -> EncryptMethod: return EncryptMethod.AES def _derive_key(self) -> bytes: - kdf = PBKDF2HMAC( - algorithm=hashes.SHA256(), - length=32, - salt=self.salt, - iterations=100000, - ) - return kdf.derive(self.secret_key) + # Optimization: Cache the key derivation for this instance, since it is expensive and can be reused + if self._derived_key is None: + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=32, + salt=self.salt, + iterations=100000, + ) + self._derived_key = kdf.derive(self.secret_key) + return self._derived_key async def encrypt(self, plaintext: str) -> str: try: @@ -55,7 +62,9 @@ async def decrypt(self, ciphertext: str) -> str: def _pad(self, data: bytes) -> bytes: block_size = 16 padding_length = block_size - (len(data) % block_size) - padding = bytes([padding_length] * padding_length) + # Optimized for both performance and clarity: avoid constructing list object + # (bytes(n * [x]) is slower than bytes([x]) * n) + padding = bytes([padding_length]) * padding_length return data + padding def _unpad(self, data: bytes) -> bytes: