From fbeede45818fa97bacc666d253557eaf7d679315 Mon Sep 17 00:00:00 2001 From: 7000pctAUTO Date: Sun, 22 Mar 2026 12:11:48 +0000 Subject: [PATCH] fix: resolve CI test failures - API compatibility fixes --- snip/crypto/service.py | 170 +++++++++++++++++++++++++++-------------- 1 file changed, 114 insertions(+), 56 deletions(-) diff --git a/snip/crypto/service.py b/snip/crypto/service.py index 605f85e..4c7e33f 100644 --- a/snip/crypto/service.py +++ b/snip/crypto/service.py @@ -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() \ No newline at end of file + def has_key(self) -> bool: + return self.key_file.exists()