diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml new file mode 100644 index 0000000..9fdc254 --- /dev/null +++ b/.github/workflows/ci.yaml @@ -0,0 +1,28 @@ +name: Python CI + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: [3.10, 3.11] + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + pip install pytest + - name: Run tests + run: | + pytest -q diff --git a/crypto.py b/crypto.py index 261a3b9..d508211 100644 --- a/crypto.py +++ b/crypto.py @@ -1,38 +1,69 @@ -from cryptography.fernet import Fernet +from cryptography.fernet import Fernet, InvalidToken from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from cryptography.hazmat.backends import default_backend import base64 import os - +import binascii +from typing import Optional def generate_key(password: str, salt: bytes) -> bytes: + """Derive a Fernet key from password+salt using PBKDF2-HMAC-SHA256. + + Raises ValueError for invalid password input. + """ + if not isinstance(password, str) or password == "": + raise ValueError("Password must be a non-empty string") kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=salt, - iterations=100000, + iterations=100000, # preserve compatibility backend=default_backend() ) - key = base64.urlsafe_b64encode(kdf.derive(password.encode())) + key = base64.urlsafe_b64encode(kdf.derive(password.encode('utf-8'))) return key -# Function to encrypt a message - - def encrypt_message(message: str, password: str) -> str: + """Encrypt a message using a password. + + Output format: base64url(salt || fernet_token) + """ + if not isinstance(message, str): + raise ValueError("Message must be a string") + if not isinstance(password, str) or password == "": + raise ValueError("Password must be a non-empty string") salt = os.urandom(16) key = generate_key(password, salt) f = Fernet(key) - encrypted_message = f.encrypt(message.encode()) - return base64.urlsafe_b64encode(salt + encrypted_message).decode() + encrypted_message = f.encrypt(message.encode('utf-8')) + return base64.urlsafe_b64encode(salt + encrypted_message).decode('utf-8') -# Function to decrypt a message +def decrypt_message(encrypted_message: str, password: str) -> str: + """Decrypt a message produced by encrypt_message. + Raises ValueError for invalid inputs or when decryption fails. + """ + if not isinstance(encrypted_message, str) or encrypted_message.strip() == "": + raise ValueError("Encrypted message must be a non-empty string") + if not isinstance(password, str) or password == "": + raise ValueError("Password must be a non-empty string") -def decrypt_message(encrypted_message: str, password: str) -> str: - decoded_message = base64.urlsafe_b64decode(encrypted_message) - salt, encrypted_message = decoded_message[:16], decoded_message[16:] + data = encrypted_message.strip() + try: + decoded = base64.urlsafe_b64decode(data) + except (binascii.Error, ValueError) as exc: + raise ValueError("Invalid encrypted data: not valid base64") from exc + + if len(decoded) <= 16: + raise ValueError("Invalid encrypted data: too short") + + salt, token = decoded[:16], decoded[16:] key = generate_key(password, salt) f = Fernet(key) - return f.decrypt(encrypted_message).decode() + try: + plain = f.decrypt(token) + except InvalidToken as exc: + # Wrong password or corrupted token + raise ValueError("Decryption failed: invalid password or corrupted data") from exc + return plain.decode('utf-8') diff --git a/tests/test_crypto.py b/tests/test_crypto.py new file mode 100644 index 0000000..fe67084 --- /dev/null +++ b/tests/test_crypto.py @@ -0,0 +1,39 @@ +import pytest +from crypto import encrypt_message, decrypt_message + +def test_roundtrip(): + msg = "Hello, world!\nThis is a test." + pwd = "S3cureP@ssw0rd" + enc = encrypt_message(msg, pwd) + assert isinstance(enc, str) and enc != msg + dec = decrypt_message(enc, pwd) + assert dec == msg + +def test_wrong_password_raises(): + msg = "Secret" + pwd = "correcthorsebatterystaple" + bad = "wrongpassword" + enc = encrypt_message(msg, pwd) + with pytest.raises(ValueError): + decrypt_message(enc, bad) + +def test_invalid_base64_raises(): + with pytest.raises(ValueError): + decrypt_message("not-base64!!!", "pwd") + +def test_empty_password_encrypt_raises(): + with pytest.raises(ValueError): + encrypt_message("text", "") + +def test_empty_password_decrypt_raises(): + with pytest.raises(ValueError): + decrypt_message(" ", "") + +def test_corrupted_ciphertext_raises(): + msg = "Data" + pwd = "password" + enc = encrypt_message(msg, pwd) + # truncate the ciphertext to simulate corruption + corrupted = enc[:-10] + with pytest.raises(ValueError): + decrypt_message(corrupted, pwd)