diff --git a/README.md b/README.md index 03030f229..a47814229 100644 --- a/README.md +++ b/README.md @@ -48,3 +48,20 @@ source .venv/bin/activate pip install -r requirements.txt streamlit run Chatbot.py ``` + + +## Karla's part + +```bash +cd "/Users/karla/Desktop/Soft Eng/llm-examples" && rm -rf .venv && /opt/homebrew/bin/python3.11 -m venv .venv && "./.venv/bin/python" -m pip install --upgrade pip setuptools wheel && "./.venv/bin/python" -m pip install streamlit cryptography pytest -q +``` + +```bash +cd "/Users/karla/Desktop/Soft Eng/llm-examples" && "./.venv/bin/python" -m pip install streamlit cryptography pytest -q +``` + +```bash +cd "/Users/karla/Desktop/Soft Eng/llm-examples" && "./.venv/bin/python" -m pytest tests/test_auth_encrypt.py -vv +``` + +### or instead of -vv use -q diff --git a/auth_encrypt.py b/auth_encrypt.py new file mode 100644 index 000000000..dd3634e1d --- /dev/null +++ b/auth_encrypt.py @@ -0,0 +1,108 @@ +""" +Simple history encryption helper using password-based key derivation. + +Public API: +- create_decrypt_key(password, salt=None): Generate salt and derive encryption key from password +- encrypt_history(plaintext, password, salt): Encrypt bytes using AES-GCM with PBKDF2-derived key +- decrypt_history(ciphertext, password, salt): Decrypt bytes using AES-GCM with PBKDF2-derived key + +""" +from __future__ import annotations + +import base64 +import os +from typing import Tuple + +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + + +def create_decrypt_key(password: str, salt: bytes = None) -> Tuple[bytes, bytes]: + """Generate a random salt and derive a 32-byte encryption key from password using PBKDF2-HMAC-SHA256. + + Args: + password: User password + salt: Optional salt bytes (16 bytes). If None, generates random salt. + + Returns: + (salt, key) tuple where: + - salt: 16-byte random salt (store this with encrypted data) + - key: 32-byte derived key for encryption/decryption + """ + if salt is None: + salt = os.urandom(16) + + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=32, + salt=salt, + iterations=200_000, + ) + key = kdf.derive(password.encode("utf-8")) + return salt, key + + +def encrypt_history(plaintext: bytes, password: str, salt: bytes) -> bytes: + """Encrypt plaintext using AES-GCM with a key derived from password and salt. + + Args: + plaintext: Data to encrypt + password: User password + salt: 16-byte salt (from create_decrypt_key) + + Returns: + Encrypted bytes in format: nonce (12 bytes) + ciphertext + """ + _, key = create_decrypt_key(password, salt) + aes = AESGCM(key) + nonce = os.urandom(12) + ciphertext = aes.encrypt(nonce, plaintext, associated_data=None) + return nonce + ciphertext + + +def decrypt_history(encrypted_data: bytes, password: str, salt: bytes) -> bytes: + """Decrypt data encrypted with encrypt_history. + + Args: + encrypted_data: Encrypted bytes (nonce + ciphertext) + password: User password + salt: 16-byte salt used during encryption + + Returns: + Decrypted plaintext bytes + + Raises: + cryptography.exceptions.InvalidTag: If password/salt incorrect or data corrupted + """ + _, key = create_decrypt_key(password, salt) + nonce = encrypted_data[:12] + ciphertext = encrypted_data[12:] + aes = AESGCM(key) + plaintext = aes.decrypt(nonce, ciphertext, associated_data=None) + return plaintext + + +if __name__ == "__main__": + # Quick local demo + demo_pw = "s3cret!" + demo_text = b"Hello! This is a secret document." + + # Create key with random salt + salt, key = create_decrypt_key(demo_pw) + print(f"Generated salt: {base64.b64encode(salt).decode()}") + + # Encrypt + encrypted = encrypt_history(demo_text, demo_pw, salt) + print(f"Encrypted {len(encrypted)} bytes") + + # Decrypt + decrypted = decrypt_history(encrypted, demo_pw, salt) + print(f"Decrypted: {decrypted.decode()}") + + # Verify wrong password fails + try: + decrypt_history(encrypted, "wrong-password", salt) + print("ERROR: Should have failed!") + except Exception as e: + print(f"Wrong password correctly rejected: {type(e).__name__}") diff --git a/auth_encrypt.py.bak b/auth_encrypt.py.bak new file mode 100644 index 000000000..839f306a2 --- /dev/null +++ b/auth_encrypt.py.bak @@ -0,0 +1,166 @@ +""" +Simple auth + file encryption helper. + +Features: +- Create user: generates salt, derives a key from password (PBKDF2-HMAC-SHA256), stores verifier and encrypts provided bytes with AES-GCM. +- Verify user: derives key and compares against stored verifier. +- Decrypt user file: locate encrypted file and decrypt with derived key. + +Storage: +- users.json in repository root stores entries like { + username: { salt: base64, kdf_iterations: int, pw_verifier: base64, enc_filename: str } +} +""" +from __future__ import annotations + +import base64 +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Tuple + +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +# Config +_ROOT = Path(__file__).resolve().parent +_USERS_FILE = _ROOT / "users.json" +_USER_DATA_DIR = _ROOT / "user_data" +_USER_DATA_DIR.mkdir(exist_ok=True) + + +def _load_users() -> dict: + if not _USERS_FILE.exists(): + return {} + with _USERS_FILE.open("r", encoding="utf-8") as f: + return json.load(f) + + +def _save_users(data: dict) -> None: + with _USERS_FILE.open("w", encoding="utf-8") as f: + json.dump(data, f, indent=2) + + +def _derive_key(password: str, salt: bytes, iterations: int = 200_000) -> bytes: + """Derive a 32-byte key from password and salt using PBKDF2-HMAC-SHA256.""" + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=32, + salt=salt, + iterations=iterations, + ) + return kdf.derive(password.encode("utf-8")) + + +def _verify_key(password: str, salt: bytes, expected: bytes, iterations: int = 200_000) -> bool: + """Return True if password derives to expected key. Uses exception semantics from KDF verify via try/except.""" + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + length=32, + salt=salt, + iterations=iterations, + ) + try: + kdf.verify(password.encode("utf-8"), expected) + return True + except Exception: + return False + + +def create_user(username: str, password: str, plaintext: bytes) -> Tuple[bool, str]: + """Create a new user entry and encrypt provided plaintext. + + Returns (success, message). + If username exists, returns False. + """ + users = _load_users() + if username in users: + return False, "username already exists" + + salt = os.urandom(16) + iterations = 200_000 + key = _derive_key(password, salt, iterations=iterations) + + # encrypt + aes = AESGCM(key) + nonce = os.urandom(12) + ct = aes.encrypt(nonce, plaintext, associated_data=None) + + enc_filename = f"{username}.enc" + enc_path = _USER_DATA_DIR / enc_filename + # store as base64: nonce + ciphertext + with enc_path.open("wb") as f: + f.write(base64.b64encode(nonce + ct)) + + users[username] = { + "salt": base64.b64encode(salt).decode("utf-8"), + "kdf_iterations": iterations, + "pw_verifier": base64.b64encode(key).decode("utf-8"), + "enc_filename": enc_filename, + } + _save_users(users) + return True, f"user created, encrypted file saved to user_data/{enc_filename}" + + +def verify_user(username: str, password: str) -> Tuple[bool, Optional[bytes]]: + """Verify username/password. If ok, returns (True, key) where key is the derived key used for encryption/decryption. + If failure, (False, None). + """ + users = _load_users() + info = users.get(username) + if not info: + return False, None + salt = base64.b64decode(info["salt"]) + iterations = int(info.get("kdf_iterations", 200_000)) + expected = base64.b64decode(info["pw_verifier"]) + # verify + ok = _verify_key(password, salt, expected, iterations=iterations) + if not ok: + return False, None + key = _derive_key(password, salt, iterations=iterations) + return True, key + + +def decrypt_user_file(username: str, key: bytes) -> Tuple[bool, Optional[bytes], str]: + """Attempt to locate and decrypt the user's encrypted file. + + Returns (success, plaintext_bytes_or_None, message). + """ + users = _load_users() + info = users.get(username) + if not info: + return False, None, "user not found" + enc_filename = info.get("enc_filename") + if not enc_filename: + return False, None, "no encrypted file for user" + enc_path = _USER_DATA_DIR / enc_filename + if not enc_path.exists(): + return False, None, f"encrypted file {enc_filename} missing" + data = base64.b64decode(enc_path.read_bytes()) + nonce = data[:12] + ct = data[12:] + aes = AESGCM(key) + try: + pt = aes.decrypt(nonce, ct, associated_data=None) + return True, pt, "decrypted" + except Exception as e: + return False, None, f"decryption failed: {e}" + + +if __name__ == "__main__": + # Quick local demo + demo_user = "alice" + demo_pw = "s3cret!" + demo_text = b"Hello Alice! This is a secret document." + + ok, msg = create_user(demo_user, demo_pw, demo_text) + print(ok, msg) + ok, key = verify_user(demo_user, demo_pw) + print("verify", ok) + if ok: + ok2, pt, m = decrypt_user_file(demo_user, key) + print(ok2, m) + if ok2: + print(pt.decode()) diff --git a/pages/6_Login_Encrypt.py b/pages/6_Login_Encrypt.py new file mode 100644 index 000000000..e78ad9636 --- /dev/null +++ b/pages/6_Login_Encrypt.py @@ -0,0 +1,55 @@ +import base64 +from pathlib import Path + +import streamlit as st + +from auth_encrypt import create_user, verify_user, decrypt_user_file, _USER_DATA_DIR + + +st.set_page_config(page_title="Login + Encrypt Demo") + +st.title("🔒 Login and File Encryption Demo") + +mode = st.radio("Mode", ["Login", "Create account"]) + +username = st.text_input("Username") +password = st.text_input("Password", type="password") + +if mode == "Create account": + uploaded = st.file_uploader("Upload a file to encrypt for this user", type=None) + if st.button("Create account and encrypt"): + if not username or not password or not uploaded: + st.error("Provide username, password and a file to encrypt") + else: + data = uploaded.read() + ok, msg = create_user(username, password, data) + if ok: + st.success(msg) + # offer download of encrypted file + enc_path = _USER_DATA_DIR / f"{username}.enc" + if enc_path.exists(): + enc_bytes = enc_path.read_bytes() + st.download_button("Download encrypted file", enc_bytes, file_name=f"{username}.enc") + else: + st.error(msg) + +else: + if st.button("Login and decrypt user's stored file"): + if not username or not password: + st.error("Provide username and password") + else: + ok, key = verify_user(username, password) + if not ok: + st.error("Invalid username or password") + else: + ok2, pt, m = decrypt_user_file(username, key) + if not ok2: + st.error(m) + else: + # try to render as text if possible + try: + text = pt.decode("utf-8") + st.text_area("Decrypted content", value=text, height=300) + except Exception: + st.success("File decrypted — binary data") + st.download_button("Download decrypted file", pt, file_name=f"{username}.decrypted") diff --git a/requirements.txt b/requirements.txt index 36459684e..04630d8ed 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ anthropic>=0.3.0 trubrics>=1.4.3 streamlit-feedback langchain-community +cryptography>=40.0.2 diff --git a/tests/test_auth_encrypt.py b/tests/test_auth_encrypt.py new file mode 100644 index 000000000..20a2fd879 --- /dev/null +++ b/tests/test_auth_encrypt.py @@ -0,0 +1,69 @@ +import pytest +from cryptography.exceptions import InvalidTag + +import auth_encrypt + + +def test_create_decrypt_key(): + """Test salt generation and key derivation.""" + password = "test-pass-123" + + # Generate new salt and key + salt1, key1 = auth_encrypt.create_decrypt_key(password) + assert len(salt1) == 16 + assert len(key1) == 32 + + # Same password with different salt produces different key + salt2, key2 = auth_encrypt.create_decrypt_key(password) + assert salt1 != salt2 + assert key1 != key2 + + # Same password with same salt produces same key + _, key3 = auth_encrypt.create_decrypt_key(password, salt=salt1) + assert key1 == key3 + + +def test_encrypt_decrypt_cycle(): + """Test full encrypt/decrypt cycle with correct password.""" + password = "s3cret-password" + plaintext = b"Secret data for pytest" + + # Generate salt and encrypt + salt, _ = auth_encrypt.create_decrypt_key(password) + encrypted = auth_encrypt.encrypt_history(plaintext, password, salt) + + # Verify encrypted data is longer (nonce + ciphertext + auth tag) + assert len(encrypted) > len(plaintext) + assert encrypted != plaintext + + # Decrypt with correct password + decrypted = auth_encrypt.decrypt_history(encrypted, password, salt) + assert decrypted == plaintext + + +def test_wrong_password_fails(): + """Test that wrong password raises exception during decryption.""" + password = "correct-password" + wrong_password = "wrong-password" + plaintext = b"Secret message" + + salt, _ = auth_encrypt.create_decrypt_key(password) + encrypted = auth_encrypt.encrypt_history(plaintext, password, salt) + + # Wrong password should raise InvalidTag + with pytest.raises(InvalidTag): + auth_encrypt.decrypt_history(encrypted, wrong_password, salt) + + +def test_wrong_salt_fails(): + """Test that wrong salt raises exception during decryption.""" + password = "password123" + plaintext = b"Another secret" + + salt, _ = auth_encrypt.create_decrypt_key(password) + encrypted = auth_encrypt.encrypt_history(plaintext, password, salt) + + # Different salt should fail + wrong_salt, _ = auth_encrypt.create_decrypt_key(password) + with pytest.raises(InvalidTag): + auth_encrypt.decrypt_history(encrypted, password, wrong_salt) diff --git a/user_data/alice.enc b/user_data/alice.enc new file mode 100644 index 000000000..9842829bc --- /dev/null +++ b/user_data/alice.enc @@ -0,0 +1 @@ +3o3Fgvrt3vgNaYjUvDPqufZvu3a/HXJ5W7kU4DTLK0cjIJXYptHHhML01lLb2D1vRM9VOpDBoBor9iY51KpT0oEiMA== \ No newline at end of file diff --git a/users.json b/users.json new file mode 100644 index 000000000..16aee35f8 --- /dev/null +++ b/users.json @@ -0,0 +1,8 @@ +{ + "alice": { + "salt": "hSt0WBARDoVPVNBnOKc6xw==", + "kdf_iterations": 200000, + "pw_verifier": "9vP6YOXKzFWrhCgFVqfk7crFKbbAtx4+UeN6Iu1emPs=", + "enc_filename": "alice.enc" + } +} \ No newline at end of file