A lightweight Python script for securely encrypting and decrypting text using AES-256-GCM.

vault.py

import os
import sys
import base64
import argparse
from pathlib import Path
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


def encrypt(plaintext: str) -> str:
    """
    Encrypt the given plaintext string using AES-256-GCM.
    
    Args:
        plaintext: The plaintext string to encrypt.
    
    Returns:
        A base64-encoded string containing the nonce concatenated with the ciphertext.
    """
    key = load()
    aesgcm = AESGCM(key)
    nonce = os.urandom(12)
    ct = aesgcm.encrypt(nonce, plaintext.encode(), None)
    return base64.b64encode(nonce + ct).decode()

def decrypt(ciphertext: str) -> str:
    """
    Decrypt the given base64-encoded ciphertext string using AES-256-GCM.
    
    Args:
        ciphertext: The base64-encoded ciphertext to decrypt (includes nonce).
    
    Returns:
        The decrypted plaintext string.
    """
    try:
        encrypted = base64.b64decode(ciphertext)
        key = load()
        aesgcm = AESGCM(key)
        nonce = encrypted[:12]
        ct = encrypted[12:]
        return aesgcm.decrypt(nonce, ct, None).decode()
    except Exception:
        raise SystemExit("Error: The text to decrypt is invalid.")

def load() -> bytes:
    """
    Load the AES key from ~/.vault.key file.
    If the file does not exist, generate and save a new key automatically.
    """
    key_path = Path.home() / ".vault.key"

    if not key_path.exists():
        key = AESGCM.generate_key(bit_length=256)
        with open(key_path, "w") as f:
            f.write(base64.b64encode(key).decode())
        return key

    with open(key_path, "r") as f:
        return base64.b64decode(f.read())
    
def main():
    """
    Parse command-line arguments and perform the requested command:
    - encrypt: Encrypt the provided text or piped input.
    - decrypt: Decrypt the provided text or piped input.
    """
    parser = argparse.ArgumentParser(description="A simple and secure AES-256-GCM encryption CLI")
    parser.add_argument("command", choices=["encrypt", "decrypt"], help="Command to run")
    parser.add_argument("text", nargs="?", help="Text to encrypt/decrypt")
    args = parser.parse_args()
    text = args.text or (sys.stdin.read().strip() if not sys.stdin.isatty() else None)

    if not text:
        raise SystemExit("Error: No input text provided. Please specify text to encrypt or decrypt.")

    if args.command == "encrypt":
        print(encrypt(text))
    elif args.command == "decrypt":
        print(decrypt(text))

if __name__ == "__main__":
    main()

# ------------------------------------------------------------------------------------------------------
# vault.py – A lightweight Python script for securely encrypting and decrypting text using AES-256-GCM.
#
# This script provides a command-line interface for:
#   1. Encrypting plaintext
#   2. Decrypting ciphertext
#
# ------------------------------------------------------------------------------------------------------
# 📦 Requirements:
#
#   • Python 3.6+
#   • Install the `cryptography` library:
#       pip install cryptography
#
# ------------------------------------------------------------------------------------------------------
# 📘 Usage:
#
#   • Encrypt plaintext:
#       python3 vault.py encrypt 'your message here'
#
#   • Decrypt ciphertext:
#       python3 vault.py decrypt 'ciphertext'
#
# ------------------------------------------------------------------------------------------------------
# 📂 File examples:
#
#   • Encrypt the contents of a file and save encrypted output:
#       python3 vault.py encrypt < input.txt > encrypted.txt
#
#   • Decrypt the contents of an encrypted file and save plaintext output:
#       python3 vault.py decrypt < encrypted.txt > decrypted.txt
#
# ------------------------------------------------------------------------------------------------------
# 💡 Pipe-based examples:
#
#   • Encrypt and decrypt a message in one line:
#       python3 vault.py encrypt 'hello world' | python3 vault.py decrypt
#
#   • Encrypt from stdin and decrypt from stdin:
#       echo 'hello world' | python3 vault.py encrypt | python3 vault.py decrypt
#
#   • Encrypt contents of a file and decrypt:
#       python3 vault.py encrypt < input.txt | python3 vault.py decrypt
#
# ------------------------------------------------------------------------------------------------------