from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.backends import default_backend
import os
import sys

def encrypt(content: bytes, password: str) -> bytes:
    """
    Encrypts the given content using the provided password.

    Uses AES-256-GCM with a PBKDF2-HMAC-SHA512 derived key, matching the
    server/browser implementation in src/crypto.ts.

    Args:
    content (bytes): The content to be encrypted.
    password (str): The password used for encryption.

    Returns:
    bytes: salt[16] | iv[12] | ciphertext+tag
    """

    # Generate a random initialization vector (IV) and salt
    iv = os.urandom(12)
    salt = os.urandom(16)

    # Derive a key from the password using PBKDF2
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA512(),
        length=32,
        salt=salt,
        iterations=210000,
        backend=default_backend()
    )
    key = kdf.derive(password.encode())

    # Encrypt with AES-256-GCM (the authentication tag is appended to the ciphertext)
    encrypted_content = AESGCM(key).encrypt(iv, content, None)

    # Return the salt, IV, and encrypted content concatenated
    return salt + iv + encrypted_content
if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: encrypt.py <file> <password> <outfile>")
        sys.exit(1)
    data = encrypt(open(sys.argv[1], "rb").read(), sys.argv[2])
    open(sys.argv[3], "wb").write(data)