image_crypt.py

#!/usr/bin/env python3
import argparse
import os
import struct
import json
from cryptography.fernet import Fernet
import pyperclip

def generate_key():
    key = Fernet.generate_key()
    pyperclip.copy(key.decode())
    print("\n🔑 Key generated and copied to clipboard!")
    print("💡 Save it securely — you'll need it later to decrypt your files.\n")
    return key

def safe_filename(filepath):
    if not os.path.exists(filepath):
        return filepath
    base, ext = os.path.splitext(filepath)
    counter = 1
    while True:
        new_name = f"{base}({counter}){ext}"
        if not os.path.exists(new_name):
            return new_name
        counter += 1

def pack_and_encrypt_file(input_path, key):
    f = Fernet(key)
    with open(input_path, "rb") as rf:
        data = rf.read()

    meta = {
        "name": os.path.basename(input_path),
        "ext": os.path.splitext(input_path)[1].lstrip('.').lower()
    }
    meta_bytes = json.dumps(meta, ensure_ascii=False).encode('utf-8')
    blob = struct.pack(">I", len(meta_bytes)) + meta_bytes + data
    encrypted = f.encrypt(blob)
    return encrypted

def decrypt_and_unpack_file(encrypted_bytes, key, output_hint=None):
    f = Fernet(key)
    blob = f.decrypt(encrypted_bytes)  # integrity checked
    meta_len = struct.unpack(">I", blob[:4])[0]
    meta_json = blob[4:4+meta_len]
    data = blob[4+meta_len:]
    meta = json.loads(meta_json.decode('utf-8', errors='replace'))

    out_name = output_hint or meta.get("name") or "output"
    if not os.path.splitext(out_name)[1] and meta.get("ext"):
        out_name = out_name + "." + meta["ext"]
    out_name = safe_filename(out_name)

    with open(out_name, "wb") as wf:
        wf.write(data)
    return out_name

def encrypt_file_cli(input_file, output_file, key):
    encrypted = pack_and_encrypt_file(input_file, key)
    with open(output_file, "wb") as wf:
        wf.write(encrypted)
    print(f"\n🔐 File '{input_file}' encrypted as '{output_file}'\n")

def decrypt_file_cli(input_file, output_file_hint, key):
    with open(input_file, "rb") as rf:
        encrypted_bytes = rf.read()
    out_name = decrypt_and_unpack_file(encrypted_bytes, key, output_file_hint)
    print(f"\n🔓 File decrypted successfully! Saved as: {out_name}\n")

def main():
    parser = argparse.ArgumentParser(description="Encrypt/decrypt files with metadata inside payload")
    parser.add_argument("action", choices=["encrypt", "decrypt", "generate_key"])
    parser.add_argument("-i", "--input")
    parser.add_argument("-o", "--output", help="Optional output filename (decrypt prefers metadata)")
    parser.add_argument("-k", "--key")
    args = parser.parse_args()

    if args.action == "generate_key":
        generate_key()
        return

    if not args.input or not args.key:
        print(f"⚠️ {args.action.capitalize()} requires --input and --key.")
        return

    key_bytes = args.key.encode()
    if args.action == "encrypt":
        out = args.output or (args.input + "_enc")
        out = safe_filename(out)
        encrypt_file_cli(args.input, out, key_bytes)
    else:
        decrypt_file_cli(args.input, args.output, key_bytes)

if __name__ == "__main__":
    main()