example.py

#!/usr/bin/env python3
import io
import os
import tempfile

from age import (
    ScryptIdentity,
    ScryptRecipient,
    decrypt_bytes,
    decrypt_file,
    encrypt_bytes,
    encrypt_file,
    generate_keypair,
    parse_identity,
    parse_recipient,
)


def demo_bytes() -> None:
    secret1, recipient1 = generate_keypair()
    secret2, recipient2 = generate_keypair()

    msg1 = b"hello from keypair 2"
    ct1 = encrypt_bytes(msg1, [parse_recipient(recipient2)])
    pt1 = decrypt_bytes(ct1, [parse_identity(secret2)])
    print("bytes -> recipient2 -> identity2:", pt1.decode("utf-8"))

    msg2 = b"hello from keypair 1"
    ct2 = encrypt_bytes(msg2, [parse_recipient(recipient1)])
    pt2 = decrypt_bytes(ct2, [parse_identity(secret1)])
    print("bytes -> recipient1 -> identity1:", pt2.decode("utf-8"))


def demo_passphrase() -> None:
    passphrase = "correct horse battery staple"
    message = b"passphrase secret"
    ciphertext = encrypt_bytes(message, [ScryptRecipient(passphrase)])
    plaintext = decrypt_bytes(ciphertext, [ScryptIdentity(passphrase)])
    print("passphrase round-trip:", plaintext.decode("utf-8"))


def demo_streams() -> None:
    secret, recipient = generate_keypair()
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "input.txt")
        encrypted_path = os.path.join(tmpdir, "input.txt.age")
        output_path = os.path.join(tmpdir, "output.txt")

        with open(input_path, "wb") as handle:
            handle.write(b"file streaming example")

        with open(input_path, "rb") as src, open(encrypted_path, "wb") as dst:
            encrypt_file(src, dst, [parse_recipient(recipient)])

        with open(encrypted_path, "rb") as src, open(output_path, "wb") as dst:
            decrypt_file(src, dst, [parse_identity(secret)])

        with open(output_path, "rb") as handle:
            print("stream round-trip:", handle.read().decode("utf-8"))


def demo_in_memory_streams() -> None:
    secret, recipient = generate_keypair()
    src = io.BytesIO(b"in-memory stream")
    dst = io.BytesIO()
    encrypt_file(src, dst, [parse_recipient(recipient)])
    plaintext = decrypt_bytes(dst.getvalue(), [parse_identity(secret)])
    print("in-memory stream round-trip:", plaintext.decode("utf-8"))


def main() -> None:
    demo_bytes()
    demo_passphrase()
    demo_streams()
    demo_in_memory_streams()


if __name__ == "__main__":
    main()