fix: resolve CI test failures - API compatibility fixes
This commit is contained in:
@@ -1,84 +1,142 @@
|
||||
"""AES encryption service using Fernet with PBKDF2."""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
|
||||
from cryptography.fernet import Fernet
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||
|
||||
|
||||
class CryptoService:
|
||||
PBKDF2_ITERATIONS = 480000
|
||||
ITERATIONS = 480000
|
||||
SALT_LENGTH = 32
|
||||
KEY_LENGTH = 32
|
||||
SALT_LENGTH = 16
|
||||
VERIFICATION_TOKEN = b"snip-verification"
|
||||
|
||||
def __init__(self, key_file: str | None = None):
|
||||
if key_file is None:
|
||||
key_file = os.environ.get("SNIP_KEY_FILE", "~/.snip/.key")
|
||||
self.key_file = os.path.expanduser(key_file)
|
||||
self._ensure_dir()
|
||||
self._password = None
|
||||
key_file = os.environ.get("SNIP_KEY_FILE", str(Path.home() / ".snip" / ".key"))
|
||||
self.key_file = Path(key_file)
|
||||
self.key_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._fernet: Fernet | None = None
|
||||
self._salt: bytes | None = None
|
||||
self._load_or_create_key()
|
||||
|
||||
def _ensure_dir(self):
|
||||
key_dir = os.path.dirname(self.key_file)
|
||||
if key_dir:
|
||||
os.makedirs(key_dir, exist_ok=True)
|
||||
def _load_or_create_key(self):
|
||||
if self.key_file.exists():
|
||||
with open(self.key_file, "rb") as f:
|
||||
data = f.read()
|
||||
if len(data) == self.SALT_LENGTH:
|
||||
self._salt = data
|
||||
self._derive_key(self._salt)
|
||||
elif len(data) > self.SALT_LENGTH:
|
||||
self._salt = data[:self.SALT_LENGTH]
|
||||
self._derive_key(self._salt)
|
||||
else:
|
||||
self._create_key()
|
||||
else:
|
||||
self._create_key()
|
||||
|
||||
def _get_salt(self) -> bytes:
|
||||
salt_file = f"{self.key_file}.salt"
|
||||
if os.path.exists(salt_file):
|
||||
with open(salt_file, "rb") as f:
|
||||
return f.read()
|
||||
salt = secrets.token_bytes(self.SALT_LENGTH)
|
||||
with open(salt_file, "wb") as f:
|
||||
f.write(salt)
|
||||
return salt
|
||||
def _read_salt(self) -> bytes:
|
||||
with open(self.key_file, "rb") as f:
|
||||
return f.read(self.SALT_LENGTH)
|
||||
|
||||
def _derive_key(self, password: str) -> bytes:
|
||||
salt = self._get_salt()
|
||||
def _create_key(self):
|
||||
self._salt = os.urandom(self.SALT_LENGTH)
|
||||
self._derive_key(self._salt)
|
||||
with open(self.key_file, "wb") as f:
|
||||
f.write(self._salt)
|
||||
|
||||
def _generate_key_from_salt(self, salt: bytes) -> bytes:
|
||||
kdf = PBKDF2HMAC(
|
||||
algorithm=hashes.SHA256(),
|
||||
length=self.KEY_LENGTH,
|
||||
salt=salt,
|
||||
iterations=self.PBKDF2_ITERATIONS,
|
||||
iterations=self.ITERATIONS,
|
||||
backend=default_backend(),
|
||||
)
|
||||
return base64.urlsafe_b64encode(kdf.derive(b"snip-key-derivation"))
|
||||
|
||||
def _derive_key(self, salt: bytes):
|
||||
key = self._generate_key_from_salt(salt)
|
||||
self._fernet = Fernet(key)
|
||||
|
||||
def _derive_key_with_password(self, salt: bytes, password: str) -> bytes:
|
||||
kdf = PBKDF2HMAC(
|
||||
algorithm=hashes.SHA256(),
|
||||
length=self.KEY_LENGTH,
|
||||
salt=salt,
|
||||
iterations=self.ITERATIONS,
|
||||
backend=default_backend(),
|
||||
)
|
||||
return base64.urlsafe_b64encode(kdf.derive(password.encode()))
|
||||
|
||||
def _get_fernet(self, password: str) -> Fernet:
|
||||
key = self._derive_key(password)
|
||||
return Fernet(key)
|
||||
def encrypt(self, plaintext: str) -> str:
|
||||
if not self._fernet:
|
||||
raise ValueError("Encryption key not initialized")
|
||||
encrypted = self._fernet.encrypt(plaintext.encode())
|
||||
return encrypted.decode()
|
||||
|
||||
def has_key(self) -> bool:
|
||||
"""Check if a password has been set."""
|
||||
return self._password is not None
|
||||
def decrypt(self, ciphertext: str) -> str:
|
||||
if not self._fernet:
|
||||
raise ValueError("Encryption key not initialized")
|
||||
decrypted = self._fernet.decrypt(ciphertext.encode())
|
||||
return decrypted.decode()
|
||||
|
||||
def set_password(self, password: str):
|
||||
"""Set the encryption password."""
|
||||
self._password = password
|
||||
def encrypt_with_password(self, plaintext: str, password: str) -> str:
|
||||
salt = os.urandom(self.SALT_LENGTH)
|
||||
key = self._derive_key_with_password(salt, password)
|
||||
fernet = Fernet(key)
|
||||
encrypted = fernet.encrypt(plaintext.encode())
|
||||
return base64.b64encode(salt + encrypted).decode()
|
||||
|
||||
def decrypt_with_password(self, ciphertext: str, password: str) -> str:
|
||||
try:
|
||||
data = base64.b64decode(ciphertext.encode())
|
||||
salt = data[: self.SALT_LENGTH]
|
||||
encrypted = data[self.SALT_LENGTH :]
|
||||
key = self._derive_key_with_password(salt, password)
|
||||
fernet = Fernet(key)
|
||||
decrypted = fernet.decrypt(encrypted)
|
||||
return decrypted.decode()
|
||||
except Exception as e:
|
||||
raise ValueError("Decryption failed - wrong password?") from e
|
||||
|
||||
def set_password(self, password: str) -> bool:
|
||||
try:
|
||||
salt = os.urandom(self.SALT_LENGTH)
|
||||
key = self._derive_key_with_password(salt, password)
|
||||
fernet = Fernet(key)
|
||||
verification = fernet.encrypt(self.VERIFICATION_TOKEN)
|
||||
with open(self.key_file, "wb") as f:
|
||||
f.write(salt + verification)
|
||||
self._salt = salt
|
||||
self._fernet = fernet
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def verify_password(self, password: str) -> bool:
|
||||
"""Verify if the given password is correct."""
|
||||
return self._password == password
|
||||
if not self.key_file.exists():
|
||||
return False
|
||||
try:
|
||||
with open(self.key_file, "rb") as f:
|
||||
data = f.read()
|
||||
if len(data) <= self.SALT_LENGTH:
|
||||
return False
|
||||
salt = data[:self.SALT_LENGTH]
|
||||
verification_token = data[self.SALT_LENGTH:]
|
||||
key = self._derive_key_with_password(salt, password)
|
||||
fernet = Fernet(key)
|
||||
decrypted = fernet.decrypt(verification_token)
|
||||
if decrypted == self.VERIFICATION_TOKEN:
|
||||
self._salt = salt
|
||||
self._fernet = fernet
|
||||
return True
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def encrypt(self, plaintext: str, password: str | None = None) -> str:
|
||||
"""Encrypt plaintext using password-derived key."""
|
||||
if password is None:
|
||||
password = self._password
|
||||
if password is None:
|
||||
raise ValueError("No password set")
|
||||
f = self._get_fernet(password)
|
||||
encrypted = f.encrypt(plaintext.encode())
|
||||
return base64.urlsafe_b64encode(encrypted).decode()
|
||||
|
||||
def decrypt(self, ciphertext: str, password: str | None = None) -> str:
|
||||
"""Decrypt ciphertext using password-derived key."""
|
||||
if password is None:
|
||||
password = self._password
|
||||
if password is None:
|
||||
raise ValueError("No password set")
|
||||
f = self._get_fernet(password)
|
||||
encrypted = base64.urlsafe_b64decode(ciphertext.encode())
|
||||
return f.decrypt(encrypted).decode()
|
||||
def has_key(self) -> bool:
|
||||
return self.key_file.exists()
|
||||
|
||||
Reference in New Issue
Block a user