UI polish and hardening

.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

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')

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)