encrypt

#!/usr/bin/env python3
import argparse
import base64
import os
import secrets
import subprocess
import sys

try:
    from cryptography.fernet import Fernet, InvalidToken
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
except ImportError:
    print("cryptography package not found. Installing...")
    try:
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "cryptography"]
        )
        from cryptography.fernet import Fernet, InvalidToken
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

        print("cryptography installed successfully")
    except (subprocess.CalledProcessError, OSError, ImportError) as e:
        print(f"Failed to install cryptography: {e}")
        sys.exit(1)


def existing_file(path: str) -> str:
    if not os.path.isfile(path):
        raise argparse.ArgumentTypeError(f"File not found: {path}")
    return path

parser = argparse.ArgumentParser(description="Encrypt a file")
parser.add_argument(
    "-f", "--file", help="File to encrypt", required=True,  type=existing_file)
parser.add_argument(
    "-k", "--key", help="Key to use for encryption", required=True)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument(
    "-e", "--encrypt", help="Encrypt the file", action="store_true")
group.add_argument(
    "-d", "--decrypt", help="Decrypt the file", action="store_true")
parser.add_argument(
    "-s", "--sneaky", help="Do not append .enc to filename", action="store_true")
args = parser.parse_args()


if args.encrypt:
    # Generate random salt for this file
    salt = secrets.token_bytes(32)

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(args.key.encode()))
    f = Fernet(key)

    with open(args.file, "rb") as file:
        try:
            data = file.read()
        except (OSError, IOError) as e:
            print(f"Error reading file: {e}")
            sys.exit(1)

    encrypted = f.encrypt(data)

    # Write salt + encrypted data
    with open(args.file, "wb") as file:
        file.write(salt + encrypted)

    if not args.sneaky:
        os.rename(args.file, args.file + ".enc")

    print("encrypted")

if args.decrypt:
    with open(args.file, "rb") as file:
        data = file.read()

    # Extract salt (first 32 bytes) and encrypted data
    if len(data) < 32:
        print("Invalid encrypted file")
        sys.exit(1)

    salt = data[:32]
    encrypted_data = data[32:]

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(args.key.encode()))
    f = Fernet(key)

    try:
        decrypted = f.decrypt(encrypted_data)
    except InvalidToken:
        print("Incorrect key")
        sys.exit(1)

    with open(args.file, "wb") as file:
        file.write(decrypted)

    if args.file.endswith(".enc"):
        outpath = args.file.replace(".enc", "")
        os.rename(args.file, outpath)

    print("decrypted")