diff --git a/pysnippets/encryption/encryption.py b/pysnippets/encryption/encryption.py index 931643b..8cfdf67 100644 --- a/pysnippets/encryption/encryption.py +++ b/pysnippets/encryption/encryption.py @@ -68,16 +68,24 @@ def unpad(self, data: bytes) -> bytes: Returns: bytes: The unpadded data. + + Raises: + ValueError: If padding is invalid. """ pad_length = data[-1] + if pad_length < 1 or pad_length > AES.block_size: + raise ValueError("Invalid padding length") + if data[-pad_length:] != bytes([pad_length] * pad_length): + raise ValueError("Invalid padding bytes") return data[:-pad_length] - def encrypt(self, data: bytes) -> str: + def encrypt(self, data: bytes, salt: bytes = b"") -> str: """ Encrypts the given data using AES encryption. Args: data (bytes): The data to encrypt. + salt (bytes): Optional salt to prepend for key derivation. Returns: str: The encrypted data encoded in base64. @@ -86,22 +94,23 @@ def encrypt(self, data: bytes) -> str: iv = get_random_bytes(AES.block_size) cipher = AES.new(self.key, AES.MODE_CBC, iv) encrypted_data = cipher.encrypt(data) - return base64.b64encode(iv + encrypted_data).decode("utf-8") + return base64.b64encode(salt + iv + encrypted_data).decode("utf-8") - def decrypt(self, enc_data: str) -> str: + def decrypt(self, enc_data: str, salt_len: int = 0) -> str: """ Decrypts the given encrypted data using AES decryption. Args: enc_data (str): The encrypted data encoded in base64. + salt_len (int): Length of the salt prepended to the data. Returns: str: The decrypted data. """ enc_data = base64.b64decode(enc_data) - iv = enc_data[:AES.block_size] + iv = enc_data[salt_len:salt_len + AES.block_size] cipher = AES.new(self.key, AES.MODE_CBC, iv) - data = cipher.decrypt(enc_data[AES.block_size:]) + data = cipher.decrypt(enc_data[salt_len + AES.block_size:]) return self.unpad(data).decode("utf-8") def encrypt_file(self, file_path: str) -> EncryptionResult: @@ -193,20 +202,34 @@ def main() -> None: The main function that initializes the AESCipher and encrypts/decrypts specified files and folders. """ passphrase = input("Enter a passphrase for encryption/decryption: ") - salt = get_random_bytes(16) # This should be securely stored to decrypt later - key = AESCipher.generate_key(passphrase, salt) - cipher = AESCipher(key) operation = input("Would you like to (e)ncrypt or (d)ecrypt? ").strip().lower() if operation == 'e': + salt = get_random_bytes(16) + key = AESCipher.generate_key(passphrase, salt) + cipher = AESCipher(key) file_path = input("Enter the path to the file to encrypt: ") folder_path = input("Enter the path to the folder to encrypt (or leave blank): ") cipher.encrypt_file(file_path) + # Save salt alongside encrypted data + salt_path = file_path + ".salt" + with open(salt_path, "wb") as f: + f.write(salt) + logging.info(f"Salt saved to {salt_path} — keep this file to decrypt later") if folder_path: cipher.encrypt_folder(folder_path) elif operation == 'd': file_path = input("Enter the path to the file to decrypt: ") + salt_path = file_path + ".salt" + try: + with open(salt_path, "rb") as f: + salt = f.read() + except FileNotFoundError: + logging.error(f"Salt file {salt_path} not found. Cannot decrypt without the original salt.") + return + key = AESCipher.generate_key(passphrase, salt) + cipher = AESCipher(key) folder_path = input("Enter the path to the folder to decrypt (or leave blank): ") cipher.decrypt_file(file_path) if folder_path: