Move SymmetricEncryption to corelibs_encryption module
This commit is contained in:
Clemens Schwaighofer
@@ -4,24 +4,11 @@ Will be moved to CoreLibs
|
||||
TODO: set key per encryption run
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import base64
|
||||
import hashlib
|
||||
from typing import TypedDict, cast
|
||||
from cryptography.fernet import Fernet
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||
import warnings
|
||||
from corelibs_encryption.symmetric import SymmetricEncryption as CorelibsSymmetricEncryption
|
||||
|
||||
|
||||
class PackageData(TypedDict):
|
||||
"""encryption package"""
|
||||
encrypted_data: str
|
||||
salt: str
|
||||
key_hash: str
|
||||
|
||||
|
||||
class SymmetricEncryption:
|
||||
class SymmetricEncryption(CorelibsSymmetricEncryption):
|
||||
"""
|
||||
simple encryption
|
||||
|
||||
@@ -29,124 +16,7 @@ class SymmetricEncryption:
|
||||
key from the password to decrypt
|
||||
"""
|
||||
|
||||
def __init__(self, password: str):
|
||||
if not password:
|
||||
raise ValueError("A password must be set")
|
||||
self.password = password
|
||||
self.password_hash = hashlib.sha256(password.encode('utf-8')).hexdigest()
|
||||
|
||||
def __derive_key_from_password(self, password: str, salt: bytes) -> bytes:
|
||||
_password = password.encode('utf-8')
|
||||
kdf = PBKDF2HMAC(
|
||||
algorithm=hashes.SHA256(),
|
||||
length=32,
|
||||
salt=salt,
|
||||
iterations=100000,
|
||||
)
|
||||
key = base64.urlsafe_b64encode(kdf.derive(_password))
|
||||
return key
|
||||
|
||||
def __encrypt_with_metadata(self, data: str | bytes) -> PackageData:
|
||||
"""Encrypt data and include salt if password-based"""
|
||||
# convert to bytes (for encoding)
|
||||
if isinstance(data, str):
|
||||
data = data.encode('utf-8')
|
||||
|
||||
# generate salt and key from password
|
||||
salt = os.urandom(16)
|
||||
key = self.__derive_key_from_password(self.password, salt)
|
||||
# init the cypher suit
|
||||
cipher_suite = Fernet(key)
|
||||
|
||||
encrypted_data = cipher_suite.encrypt(data)
|
||||
|
||||
# If using password, include salt in the result
|
||||
return {
|
||||
'encrypted_data': base64.urlsafe_b64encode(encrypted_data).decode('utf-8'),
|
||||
'salt': base64.urlsafe_b64encode(salt).decode('utf-8'),
|
||||
'key_hash': hashlib.sha256(key).hexdigest()
|
||||
}
|
||||
|
||||
def encrypt_with_metadata(self, data: str | bytes, return_as: str = 'str') -> str | bytes | PackageData:
|
||||
"""encrypt with metadata, but returns data in string"""
|
||||
match return_as:
|
||||
case 'str':
|
||||
return self.encrypt_with_metadata_return_str(data)
|
||||
case 'json':
|
||||
return self.encrypt_with_metadata_return_str(data)
|
||||
case 'bytes':
|
||||
return self.encrypt_with_metadata_return_bytes(data)
|
||||
case 'dict':
|
||||
return self.encrypt_with_metadata_return_dict(data)
|
||||
case _:
|
||||
# default is string json
|
||||
return self.encrypt_with_metadata_return_str(data)
|
||||
|
||||
def encrypt_with_metadata_return_dict(self, data: str | bytes) -> PackageData:
|
||||
"""encrypt with metadata, but returns data as PackageData dict"""
|
||||
return self.__encrypt_with_metadata(data)
|
||||
|
||||
def encrypt_with_metadata_return_str(self, data: str | bytes) -> str:
|
||||
"""encrypt with metadata, but returns data in string"""
|
||||
return json.dumps(self.__encrypt_with_metadata(data))
|
||||
|
||||
def encrypt_with_metadata_return_bytes(self, data: str | bytes) -> bytes:
|
||||
"""encrypt with metadata, but returns data in bytes"""
|
||||
return json.dumps(self.__encrypt_with_metadata(data)).encode('utf-8')
|
||||
|
||||
def decrypt_with_metadata(self, encrypted_package: str | bytes | PackageData, password: str | None = None) -> str:
|
||||
"""Decrypt data that may include metadata"""
|
||||
try:
|
||||
# Try to parse as JSON (password-based encryption)
|
||||
if isinstance(encrypted_package, bytes):
|
||||
package_data = cast(PackageData, json.loads(encrypted_package.decode('utf-8')))
|
||||
elif isinstance(encrypted_package, str):
|
||||
package_data = cast(PackageData, json.loads(str(encrypted_package)))
|
||||
else:
|
||||
package_data = encrypted_package
|
||||
|
||||
encrypted_data = base64.urlsafe_b64decode(package_data['encrypted_data'])
|
||||
salt = base64.urlsafe_b64decode(package_data['salt'])
|
||||
pwd = password or self.password
|
||||
key = self.__derive_key_from_password(pwd, salt)
|
||||
if package_data['key_hash'] != hashlib.sha256(key).hexdigest():
|
||||
raise ValueError("Key hash is not matching, possible invalid password")
|
||||
cipher_suite = Fernet(key)
|
||||
decrypted_data = cipher_suite.decrypt(encrypted_data)
|
||||
|
||||
except (json.JSONDecodeError, KeyError, UnicodeDecodeError) as e:
|
||||
raise ValueError(f"Invalid encrypted package format {e}") from e
|
||||
|
||||
return decrypted_data.decode('utf-8')
|
||||
|
||||
@staticmethod
|
||||
def encrypt_data(data: str | bytes, password: str) -> str:
|
||||
"""
|
||||
Static method to encrypt some data
|
||||
|
||||
Arguments:
|
||||
data {str | bytes} -- _description_
|
||||
password {str} -- _description_
|
||||
|
||||
Returns:
|
||||
str -- _description_
|
||||
"""
|
||||
encryptor = SymmetricEncryption(password)
|
||||
return encryptor.encrypt_with_metadata_return_str(data)
|
||||
|
||||
@staticmethod
|
||||
def decrypt_data(data: str | bytes | PackageData, password: str) -> str:
|
||||
"""
|
||||
Static method to decrypt some data
|
||||
|
||||
Arguments:
|
||||
data {str | bytes | PackageData} -- _description_
|
||||
password {str} -- _description_
|
||||
|
||||
Returns:
|
||||
str -- _description_
|
||||
"""
|
||||
decryptor = SymmetricEncryption(password)
|
||||
return decryptor.decrypt_with_metadata(data, password=password)
|
||||
warnings.warn("Use corelibs_encryption.symmetric.SymmetricEncryption instead", DeprecationWarning, stacklevel=2)
|
||||
|
||||
# __END__
|
||||
|
||||
Reference in New Issue
Block a user