AES-256 File Encryption Tool (python, written by Codex)

Codex wrote this python project from the task below. It installed and ran on a clean Ubuntu 22.04 machine as written.

Task given to the agent:

TASK: AES-256 File Encryption Tool

Write a program that encrypts and decrypts files using AES-256 encryption in CBC mode with proper key derivation, initialization vectors, and authenticated encryption to prevent tampering.

FUNCTIONAL REQUIREMENTS:
- Accept a file path, a password, and a mode (encrypt or decrypt) as command-line arguments
- Derive the encryption key from the password using PBKDF2 with SHA-256, a random 16-byte salt, and at least 100,000 iterations
- Encrypt using AES-256-CBC with a random 16-byte initialization vector (IV) for each encryption
- Prepend the salt and IV to the encrypted output file so they are available for decryption
- Add HMAC-SHA256 authentication tag to the encrypted file to detect tampering during decryption
- During decryption, verify the HMAC before attempting to decrypt and report a clear error if the file has been modified
- Support encrypting entire directories recursively via --recursive flag, preserving directory structure in the output
- Support a --output flag to specify the output file or directory (default: append .enc for encryption, strip .enc for decryption)
- Display progress information for large files: file size, percentage complete, and throughput
- If no input file is given, generate a sample text file with random content, encrypt it, then decrypt it back and verify the output matches the original
- Handle errors gracefully: wrong password (HMAC mismatch), corrupted file, file not found, permission denied

Create a complete Python project for a clean Ubuntu 22.04 machine with only Python 3.10+ installed. Include:
- Source code
- requirements.txt with all dependencies (direct and transitive) pinned to exact versions
- README.md with setup instructions, dependency explanations, build steps, run commands, and expected output

[]

# AES-256 File Encryption Tool (Python)

Encrypts/decrypts files and directories with:

- PBKDF2-SHA256 (`120000` iterations, random 16-byte salt)
- AES-256-CBC (random 16-byte IV)
- HMAC-SHA256 authentication tag

## Requirements

- Ubuntu 22.04
- Python 3.10+

## Dependencies

- Direct:
  - `pycryptodome==3.21.0`
- Transitive:
  - none

Pinned in `requirements.txt`.

## Setup

```bash
python -m pip install -r requirements.txt
```

## Usage

```bash
python src/main.py <input-path> <password> <encrypt|decrypt> [--recursive] [--output <path>]
```

Examples:

```bash
python src/main.py ./secret.txt myPassword encrypt
python src/main.py ./secret.txt.enc myPassword decrypt
python src/main.py ./data myPassword encrypt --recursive --output ./encrypted_data
```

No arguments runs a self-test:

```bash
python src/main.py
```

pycryptodome==3.21.0

#!/usr/bin/env python3
import argparse
import hashlib
import hmac
import os
import secrets
import time
from pathlib import Path
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


MAGIC = b"A2CB"
SALT_LEN = 16
IV_LEN = 16
HMAC_LEN = 32
PBKDF2_ITERS = 120000
CHUNK_SIZE = 1024 * 1024


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="AES-256 File Encryption Tool")
    parser.add_argument("input_path", nargs="?", help="Input file or directory")
    parser.add_argument("password", nargs="?", help="Password")
    parser.add_argument("mode", nargs="?", choices=["encrypt", "decrypt"], help="Mode")
    parser.add_argument("--recursive", action="store_true", help="Process directories recursively")
    parser.add_argument("--output", help="Output file or directory")
    return parser.parse_args()


def derive_keys(password: str, salt: bytes) -> tuple[bytes, bytes]:
    key_material = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, PBKDF2_ITERS, dklen=64)
    return key_material[:32], key_material[32:]


def print_progress(label: str, done: int, total: int, start: float) -> None:
    elapsed = max(1e-9, time.time() - start)
    pct = 100.0 if total == 0 else (done / total) * 100.0
    throughput = done / elapsed / (1024 * 1024)
    print(f"\r{label} | {total} bytes | {pct:5.1f}% | {throughput:6.2f} MiB/s", end="", flush=True)
    if done >= total:
        print("")


def encrypt_bytes(data: bytes, password: str, label: str = "Encrypting") -> bytes:
    salt = secrets.token_bytes(SALT_LEN)
    iv = secrets.token_bytes(IV_LEN)
    enc_key, mac_key = derive_keys(password, salt)
    cipher = AES.new(enc_key, AES.MODE_CBC, iv)

    encrypted_chunks = []
    start = time.time()
    processed = 0
    for i in range(0, len(data), CHUNK_SIZE):
        chunk = data[i : i + CHUNK_SIZE]
        processed += len(chunk)
        if i + CHUNK_SIZE < len(data):
            encrypted_chunks.append(cipher.encrypt(chunk))
        else:
            encrypted_chunks.append(cipher.encrypt(pad(chunk, AES.block_size)))
        print_progress(label, processed, len(data), start)
    if not data:
        encrypted_chunks.append(cipher.encrypt(pad(b"", AES.block_size)))
        print_progress(label, 0, 0, start)
    ciphertext = b"".join(encrypted_chunks)
    header = MAGIC + salt + iv
    tag = hmac.new(mac_key, header + ciphertext, hashlib.sha256).digest()
    return header + ciphertext + tag


def decrypt_bytes(data: bytes, password: str, label: str = "Decrypting") -> bytes:
    if len(data) < len(MAGIC) + SALT_LEN + IV_LEN + HMAC_LEN:
        raise ValueError("Corrupted file: too short.")
    if data[: len(MAGIC)] != MAGIC:
        raise ValueError("Corrupted file: invalid header.")

    salt_start = len(MAGIC)
    iv_start = salt_start + SALT_LEN
    body_start = iv_start + IV_LEN
    tag_start = len(data) - HMAC_LEN
    salt = data[salt_start:iv_start]
    iv = data[iv_start:body_start]
    ciphertext = data[body_start:tag_start]
    stored_tag = data[tag_start:]

    enc_key, mac_key = derive_keys(password, salt)
    calc_tag = hmac.new(mac_key, data[:tag_start], hashlib.sha256).digest()
    if not hmac.compare_digest(stored_tag, calc_tag):
        raise ValueError("HMAC mismatch: wrong password or file has been modified.")

    cipher = AES.new(enc_key, AES.MODE_CBC, iv)
    plain_chunks = []
    start = time.time()
    processed = 0
    for i in range(0, len(ciphertext), CHUNK_SIZE):
        chunk = ciphertext[i : i + CHUNK_SIZE]
        plain_chunks.append(cipher.decrypt(chunk))
        processed += len(chunk)
        print_progress(label, processed, len(ciphertext), start)
    plaintext = b"".join(plain_chunks)
    return unpad(plaintext, AES.block_size)


def default_output(input_path: Path, mode: str) -> Path:
    if mode == "encrypt":
        return Path(str(input_path) + ".enc")
    if str(input_path).endswith(".enc"):
        return Path(str(input_path)[:-4])
    return Path(str(input_path) + ".dec")


def process_file(input_path: Path, output_path: Path, password: str, mode: str) -> None:
    data = input_path.read_bytes()
    if mode == "encrypt":
        out = encrypt_bytes(data, password, label=f"Encrypting {input_path.name}")
    else:
        out = decrypt_bytes(data, password, label=f"Decrypting {input_path.name}")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(out)


def collect_files(root: Path) -> list[Path]:
    return [p for p in root.rglob("*") if p.is_file()]


def process_path(input_path: Path, output: Path | None, password: str, mode: str, recursive: bool) -> None:
    if input_path.is_file():
        out_file = output if output else default_output(input_path, mode)
        process_file(input_path, out_file, password, mode)
        print(f"{mode} complete: {input_path} -> {out_file}")
        return
    if not input_path.is_dir():
        raise ValueError("Input path must be file or directory.")
    if not recursive:
        raise ValueError("Input is a directory. Use --recursive.")
    out_root = output if output else Path(str(input_path) + f"_{mode}ed")
    for src in collect_files(input_path):
        rel = src.relative_to(input_path)
        target = default_output(out_root / rel, mode)
        process_file(src, target, password, mode)
        print(f"{mode} complete: {src} -> {target}")


def random_text(size: int) -> str:
    return secrets.token_hex(size)


def run_self_test(password: str) -> None:
    sample = Path("sample_plain.txt").resolve()
    enc = Path("sample_plain.txt.enc").resolve()
    dec = Path("sample_plain.txt.dec").resolve()
    sample.write_text(random_text(4096), encoding="utf-8")
    process_file(sample, enc, password, "encrypt")
    process_file(enc, dec, password, "decrypt")
    if sample.read_bytes() != dec.read_bytes():
        raise ValueError("Self-test failed: decrypted output mismatch.")
    print("Self-test successful:")
    print(f"  {sample}")
    print(f"  {enc}")
    print(f"  {dec}")


def main() -> int:
    args = parse_args()
    if not args.input_path:
        run_self_test("sample-password-123")
        return 0

    input_path = Path(args.input_path).resolve()
    if not input_path.exists():
        print(f"Error: file not found: {input_path}")
        return 1
    if not args.password or not args.mode:
        print("Error: provide <input_path> <password> <mode>.")
        return 1

    output = Path(args.output).resolve() if args.output else None
    try:
        process_path(input_path, output, args.password, args.mode, args.recursive)
        return 0
    except PermissionError:
        print("Error: permission denied.")
        return 1
    except ValueError as exc:
        print(f"Error: {exc}")
        return 1
    except Exception as exc:
        print(f"Error: {exc}")
        return 1


if __name__ == "__main__":
    raise SystemExit(main())