Tyler
137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
import os
|
|
import base64
|
|
from typing import Optional
|
|
|
|
from cryptography.hazmat.primitives import padding, hashes, hmac
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
from cryptography.hazmat.backends import default_backend
|
|
from cryptography.exceptions import InvalidSignature
|
|
|
|
# === Constants ===
|
|
SALT_LENGTH = 16
|
|
IV_LENGTH = 16
|
|
PBKDF2_ITERATIONS = 200_000
|
|
KEY_LENGTH = 32
|
|
HMAC_KEY_LENGTH = 32 # For HMAC-SHA256
|
|
HMAC_LENGTH = 32 # Output size of SHA256
|
|
|
|
# === Base64 Helpers ===
|
|
def b64encode(data: bytes) -> str:
|
|
return base64.b64encode(data).decode('utf-8')
|
|
|
|
def b64decode(data: str) -> bytes:
|
|
return base64.b64decode(data.encode('utf-8'))
|
|
|
|
# === Key Derivation ===
|
|
def derive_key(password: str, salt: bytes) -> bytes:
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=KEY_LENGTH + HMAC_KEY_LENGTH,
|
|
salt=salt,
|
|
iterations=PBKDF2_ITERATIONS,
|
|
backend=default_backend()
|
|
)
|
|
full_key = kdf.derive(password.encode('utf-8'))
|
|
return full_key[:KEY_LENGTH], full_key[KEY_LENGTH:]
|
|
|
|
# === Encrypt Text ===
|
|
def encrypt_text(plaintext: str, password: str) -> str:
|
|
salt = os.urandom(SALT_LENGTH)
|
|
iv = os.urandom(IV_LENGTH)
|
|
aes_key, hmac_key = derive_key(password, salt)
|
|
|
|
padder = padding.PKCS7(128).padder()
|
|
padded = padder.update(plaintext.encode('utf-8')) + padder.finalize()
|
|
|
|
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
|
encryptor = cipher.encryptor()
|
|
ciphertext = encryptor.update(padded) + encryptor.finalize()
|
|
|
|
payload = salt + iv + ciphertext
|
|
|
|
h = hmac.HMAC(hmac_key, hashes.SHA256(), backend=default_backend())
|
|
h.update(payload)
|
|
mac = h.finalize()
|
|
|
|
return b64encode(payload + mac)
|
|
|
|
# === Decrypt Text ===
|
|
def decrypt_text(encrypted_b64: str, password: str) -> str:
|
|
raw = b64decode(encrypted_b64)
|
|
|
|
salt = raw[:SALT_LENGTH]
|
|
iv = raw[SALT_LENGTH:SALT_LENGTH + IV_LENGTH]
|
|
ciphertext = raw[SALT_LENGTH + IV_LENGTH:-HMAC_LENGTH]
|
|
mac = raw[-HMAC_LENGTH:]
|
|
|
|
aes_key, hmac_key = derive_key(password, salt)
|
|
|
|
h = hmac.HMAC(hmac_key, hashes.SHA256(), backend=default_backend())
|
|
h.update(raw[:-HMAC_LENGTH])
|
|
h.verify(mac)
|
|
|
|
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
|
decryptor = cipher.decryptor()
|
|
padded = decryptor.update(ciphertext) + decryptor.finalize()
|
|
|
|
unpadder = padding.PKCS7(128).unpadder()
|
|
plaintext = unpadder.update(padded) + unpadder.finalize()
|
|
|
|
return plaintext.decode('utf-8')
|
|
|
|
# === Encrypt File ===
|
|
def encrypt_file(in_path, out_path, password: str, metadata: Optional[dict] = None):
|
|
with open(in_path, 'rb') as f:
|
|
plaintext = f.read()
|
|
|
|
salt = os.urandom(SALT_LENGTH)
|
|
iv = os.urandom(IV_LENGTH)
|
|
aes_key, hmac_key = derive_key(password, salt)
|
|
|
|
padder = padding.PKCS7(128).padder()
|
|
padded = padder.update(plaintext) + padder.finalize()
|
|
|
|
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
|
encryptor = cipher.encryptor()
|
|
ciphertext = encryptor.update(padded) + encryptor.finalize()
|
|
|
|
payload = salt + iv + ciphertext
|
|
|
|
h = hmac.HMAC(hmac_key, hashes.SHA256(), backend=default_backend())
|
|
h.update(payload)
|
|
mac = h.finalize()
|
|
|
|
with open(out_path, 'wb') as f:
|
|
f.write(payload + mac)
|
|
|
|
# === Decrypt File ===
|
|
def decrypt_file(in_path, out_path, password: str, metadata: Optional[dict] = None):
|
|
with open(in_path, 'rb') as f:
|
|
raw = f.read()
|
|
|
|
salt = raw[:SALT_LENGTH]
|
|
iv = raw[SALT_LENGTH:SALT_LENGTH + IV_LENGTH]
|
|
ciphertext = raw[SALT_LENGTH + IV_LENGTH:-HMAC_LENGTH]
|
|
mac = raw[-HMAC_LENGTH:]
|
|
|
|
aes_key, hmac_key = derive_key(password, salt)
|
|
|
|
h = hmac.HMAC(hmac_key, hashes.SHA256(), backend=default_backend())
|
|
h.update(raw[:-HMAC_LENGTH])
|
|
h.verify(mac)
|
|
|
|
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
|
decryptor = cipher.decryptor()
|
|
padded = decryptor.update(ciphertext) + decryptor.finalize()
|
|
|
|
unpadder = padding.PKCS7(128).unpadder()
|
|
plaintext = unpadder.update(padded) + unpadder.finalize()
|
|
|
|
with open(out_path, 'wb') as f:
|
|
f.write(plaintext)
|
|
|
|
# === Algo Name ===
|
|
def get_name():
|
|
return "AES-CBC"
|