95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
import os
|
|
import base64
|
|
from typing import Dict, Any, Union
|
|
|
|
from cryptography.fernet import Fernet
|
|
from cryptography.hazmat.primitives import hashes, padding
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
|
|
|
|
ENCRYPTION_PREFIX: str = "__encrypt__"
|
|
|
|
|
|
def get_encryption_key() -> bytes:
|
|
key: Union[str, None] = os.environ.get("NETMIKO_TOOLS_KEY")
|
|
if not key:
|
|
raise ValueError(
|
|
"Encryption key not found. Set the 'NETMIKO_TOOLS_KEY' environment variable."
|
|
)
|
|
return key.encode()
|
|
|
|
|
|
def decrypt_value(encrypted_value: str, key: bytes, encryption_type: str) -> str:
|
|
# Remove the encryption prefix
|
|
encrypted_value = encrypted_value.replace(ENCRYPTION_PREFIX, "", 1)
|
|
|
|
# Extract salt and ciphertext
|
|
salt_str, ciphertext_str = encrypted_value.split(":", 1)
|
|
salt = base64.b64decode(salt_str)
|
|
ciphertext = base64.b64decode(ciphertext_str)
|
|
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
)
|
|
derived_key: bytes = kdf.derive(key)
|
|
|
|
if encryption_type == "fernet":
|
|
f = Fernet(base64.urlsafe_b64encode(derived_key))
|
|
return f.decrypt(ciphertext).decode()
|
|
elif encryption_type == "aes128":
|
|
iv = ciphertext[:16]
|
|
ciphertext = ciphertext[16:]
|
|
cipher = Cipher(algorithms.AES(derived_key[:16]), modes.CBC(iv))
|
|
decryptor = cipher.decryptor()
|
|
padded: bytes = decryptor.update(ciphertext) + decryptor.finalize()
|
|
unpadder = padding.PKCS7(128).unpadder()
|
|
unpadded: bytes = unpadder.update(padded) + unpadder.finalize()
|
|
return unpadded.decode()
|
|
else:
|
|
raise ValueError(f"Unsupported encryption type: {encryption_type}")
|
|
|
|
|
|
def decrypt_config(
|
|
config: Dict[str, Any], key: bytes, encryption_type: str
|
|
) -> Dict[str, Any]:
|
|
for device, params in config.items():
|
|
if isinstance(params, dict):
|
|
for field, value in params.items():
|
|
if isinstance(value, str) and value.startswith(ENCRYPTION_PREFIX):
|
|
len_prefix = len(ENCRYPTION_PREFIX)
|
|
data: str = value[len_prefix:]
|
|
params[field] = decrypt_value(data, key, encryption_type)
|
|
return config
|
|
|
|
|
|
def encrypt_value(value: str, key: bytes, encryption_type: str) -> str:
|
|
salt = os.urandom(16)
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
)
|
|
derived_key: bytes = kdf.derive(key)
|
|
|
|
if encryption_type == "fernet":
|
|
f = Fernet(base64.urlsafe_b64encode(derived_key))
|
|
encrypted = f.encrypt(value.encode())
|
|
elif encryption_type == "aes128":
|
|
iv = os.urandom(16)
|
|
padder = padding.PKCS7(128).padder()
|
|
padded_data = padder.update(value.encode()) + padder.finalize()
|
|
cipher = Cipher(algorithms.AES(derived_key[:16]), modes.CBC(iv))
|
|
encryptor = cipher.encryptor()
|
|
encrypted = iv + encryptor.update(padded_data) + encryptor.finalize()
|
|
else:
|
|
raise ValueError(f"Unsupported encryption type: {encryption_type}")
|
|
|
|
# Combine salt and encrypted data
|
|
b64_salt = base64.b64encode(salt).decode()
|
|
return f"{ENCRYPTION_PREFIX}{b64_salt}:{base64.b64encode(encrypted).decode()}"
|