Update encryption_utils.py

iot-integration/data_encryption/encryption_utils.py

@@ -1,33 +1,42 @@
-import hashlib
+import os
 import base64
 from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
-def generate_secret_key():
-    """Generate a secret key for encryption"""
-    return Fernet.generate_key()
+def generate_key(password, salt):
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=100000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
+    return key
 
-def encrypt_data(data, secret_key):
-    """Encrypt data using a secret key"""
-    f = Fernet(secret_key)
-    encrypted_data = f.encrypt(data.encode())
-    return encrypted_data.decode()
+def encrypt_data(key, data):
+    cipher_suite = Fernet(key)
+    encrypted_data = cipher_suite.encrypt(data.encode())
+    return encrypted_data
 
-def decrypt_data(encrypted_data, secret_key):
-    """Decrypt data using a secret key"""
-    f = Fernet(secret_key)
-    decrypted_data = f.decrypt(encrypted_data.encode())
-    return decrypted_data.decode()
+def decrypt_data(key, encrypted_data):
+    cipher_suite = Fernet(key)
+    decrypted_data = cipher_suite.decrypt(encrypted_data).decode()
+    return decrypted_data
 
-def hash_data(data):
-    """Hash data using SHA-256"""
-    h = hashlib.sha256()
-    h.update(data.encode())
-    return h.hexdigest()
+def generate_salt():
+    salt = os.urandom(16)
+    return salt
 
-def base64_encode(data):
-    """Base64 encode data"""
-    return base64.b64encode(data.encode()).decode()
+def main():
+    password = 'my_secret_password'
+    salt = generate_salt()
+    key = generate_key(password, salt)
+    data = 'This is some secret data'
+    encrypted_data = encrypt_data(key, data)
+    print('Encrypted data:', encrypted_data)
+    decrypted_data = decrypt_data(key, encrypted_data)
+    print('Decrypted data:', decrypted_data)
 
-def base64_decode(encoded_data):
-    """Base64 decode data"""
-    return base64.b64decode(encoded_data.encode()).decode()
+if __name__ == '__main__':
+    main()