clio/clio-api/test_crypto.py

#!/usr/bin/env python3
"""
Unit tests for the crypto module.

Tests encryption/decryption roundtrip without needing the server running.

Usage:
    python test_crypto.py
"""

import io
import os
import struct
import tempfile
import unittest
from pathlib import Path

from nacl.public import PrivateKey, PublicKey, SealedBox

# Import the crypto module
import crypto

# SREC format constants (must match Android app)
MAGIC_BYTES = b"SREC"
VERSION = 1
CHUNK_SIZE = 64 * 1024


def encrypt_like_android(data: bytes, public_key_bytes: bytes) -> bytes:
    """
    Simulates the Android app's encryption.
    This is a reference implementation to verify the server can decrypt it.
    """
    public_key = PublicKey(public_key_bytes)
    sealed_box = SealedBox(public_key)

    output = io.BytesIO()

    # Calculate chunks
    chunks = [data[i:i + CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
    chunk_count = len(chunks)

    # Write header
    output.write(MAGIC_BYTES)
    output.write(struct.pack(">B", VERSION))
    output.write(struct.pack(">I", chunk_count))

    # Encrypt and write each chunk
    for chunk in chunks:
        encrypted = sealed_box.encrypt(chunk)
        output.write(struct.pack(">I", len(encrypted)))
        output.write(encrypted)

    return output.getvalue()


class TestCrypto(unittest.TestCase):

    def setUp(self):
        """Set up test fixtures."""
        # Create a temporary directory for test keys
        self.temp_dir = tempfile.mkdtemp()
        self.original_keys_dir = crypto.config.KEYS_DIR
        self.original_private_key_file = crypto.config.PRIVATE_KEY_FILE
        self.original_public_key_file = crypto.config.PUBLIC_KEY_FILE

        # Point to temp directory
        crypto.config.KEYS_DIR = Path(self.temp_dir)
        crypto.config.PRIVATE_KEY_FILE = Path(self.temp_dir) / "private.key"
        crypto.config.PUBLIC_KEY_FILE = Path(self.temp_dir) / "public.key"

    def tearDown(self):
        """Clean up test fixtures."""
        # Restore original paths
        crypto.config.KEYS_DIR = self.original_keys_dir
        crypto.config.PRIVATE_KEY_FILE = self.original_private_key_file
        crypto.config.PUBLIC_KEY_FILE = self.original_public_key_file

        # Clean up temp directory
        import shutil
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_keypair_generation(self):
        """Test that keypair generation works."""
        private_key, public_key = crypto.generate_keypair()

        self.assertEqual(len(private_key), 32)
        self.assertEqual(len(public_key), 32)
        self.assertNotEqual(private_key, public_key)

    def test_keypair_save_load(self):
        """Test that keypairs can be saved and loaded."""
        # Generate and save
        private_key, public_key = crypto.generate_keypair()
        crypto.save_keypair(private_key, public_key)

        # Load
        loaded_private, loaded_public = crypto.load_or_create_keypair()

        self.assertEqual(private_key, loaded_private)
        self.assertEqual(public_key, loaded_public)

    def test_public_key_base64(self):
        """Test Base64 encoding of public key."""
        import base64

        _, public_key = crypto.load_or_create_keypair()
        b64 = crypto.get_public_key_base64()

        # Verify it's valid Base64 that decodes to the key
        decoded = base64.b64decode(b64)
        self.assertEqual(decoded, public_key)

    def test_encryption_decryption_small_file(self):
        """Test roundtrip encryption/decryption of a small file."""
        # Generate keypair
        private_key, public_key = crypto.load_or_create_keypair()

        # Create test data (smaller than one chunk)
        original_data = b"Hello, this is a test recording!" * 100

        # Encrypt like Android app would
        encrypted = encrypt_like_android(original_data, public_key)

        # Write encrypted to temp file
        encrypted_path = Path(self.temp_dir) / "test.ogg.enc"
        decrypted_path = Path(self.temp_dir) / "test.ogg"

        encrypted_path.write_bytes(encrypted)

        # Decrypt using server's crypto module
        success = crypto.decrypt_file(encrypted_path, decrypted_path)

        self.assertTrue(success)
        self.assertEqual(decrypted_path.read_bytes(), original_data)

    def test_encryption_decryption_large_file(self):
        """Test roundtrip encryption/decryption of a large file (multiple chunks)."""
        # Generate keypair
        private_key, public_key = crypto.load_or_create_keypair()

        # Create test data (larger than one chunk = 64KB)
        original_data = os.urandom(200 * 1024)  # 200KB

        # Encrypt like Android app would
        encrypted = encrypt_like_android(original_data, public_key)

        # Write encrypted to temp file
        encrypted_path = Path(self.temp_dir) / "test_large.ogg.enc"
        decrypted_path = Path(self.temp_dir) / "test_large.ogg"

        encrypted_path.write_bytes(encrypted)

        # Decrypt using server's crypto module
        success = crypto.decrypt_file(encrypted_path, decrypted_path)

        self.assertTrue(success)
        self.assertEqual(decrypted_path.read_bytes(), original_data)

    def test_invalid_magic_bytes(self):
        """Test that invalid magic bytes are rejected."""
        encrypted_path = Path(self.temp_dir) / "invalid.enc"
        decrypted_path = Path(self.temp_dir) / "invalid.ogg"

        # Write file with wrong magic bytes
        encrypted_path.write_bytes(b"XXXX" + b"\x01" + b"\x00\x00\x00\x01")

        success = crypto.decrypt_file(encrypted_path, decrypted_path)
        self.assertFalse(success)

    def test_invalid_version(self):
        """Test that unsupported versions are rejected."""
        encrypted_path = Path(self.temp_dir) / "invalid_version.enc"
        decrypted_path = Path(self.temp_dir) / "invalid_version.ogg"

        # Write file with unsupported version
        encrypted_path.write_bytes(MAGIC_BYTES + b"\x99" + b"\x00\x00\x00\x01")

        success = crypto.decrypt_file(encrypted_path, decrypted_path)
        self.assertFalse(success)

    def test_output_path_generation(self):
        """Test that output paths are generated correctly."""
        # Test normal filename
        path = crypto.get_output_path_for_recording("recording_2024-01-15_09-30-00.ogg.enc")
        self.assertTrue(str(path).endswith("2024/01/15/recording_2024-01-15_09-30-00.ogg"))

        # Test another date
        path = crypto.get_output_path_for_recording("recording_2023-12-25_23-59-59.ogg.enc")
        self.assertTrue(str(path).endswith("2023/12/25/recording_2023-12-25_23-59-59.ogg"))


if __name__ == "__main__":
    unittest.main(verbosity=2)