from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
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.

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

    Returns:
    bytes: The encrypted content.
    """

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

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

    # Create a cipher object with AES-256-CBC
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()

    # Pad the content to a multiple of the block length
    padder = padding.PKCS7(128).padder()
    padded_content = padder.update(content) + padder.finalize()

    # Encrypt the padded content
    encrypted_content = encryptor.update(padded_content) + encryptor.finalize()

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