//	export PASSWORD="my secret password"
//	go run . encrypt "secret text"  > vault.bin
//	go run . decrypt < vault.bin
//
// (the password is read from the PASSWORD env var to keep the example simple)
package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/pbkdf2"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"io"
	"os"
)

const (
	magic   = "DEMO1"
	saltLen = 16      // random salt for PBKDF2
	keyLen  = 32      // AES-256
	iter    = 600_000 // PBKDF2 iterations
)

func fatal(err error) {
	fmt.Fprintln(os.Stderr, "error:", err)
	os.Exit(1)
}

// newGCM derives the key from the password (PBKDF2-SHA256) and builds AES-256-GCM.
func newGCM(pass string, salt []byte) (cipher.AEAD, error) {
	key, err := pbkdf2.Key(sha256.New, pass, salt, iter, keyLen)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

// encrypt produces: magic | salt | nonce | ciphertext+tag.
func encrypt(pass string, plain []byte) ([]byte, error) {
	salt := make([]byte, saltLen)
	if _, err := rand.Read(salt); err != nil {
		return nil, err
	}
	gcm, err := newGCM(pass, salt)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}
	ct := gcm.Seal(nil, nonce, plain, []byte(magic))

	out := append([]byte(magic), salt...)
	out = append(out, nonce...)
	return append(out, ct...), nil
}

// decrypt reverses encrypt and authenticates via the same AAD (magic).
func decrypt(pass string, blob []byte) ([]byte, error) {
	if len(blob) < len(magic)+saltLen+12 || string(blob[:len(magic)]) != magic {
		return nil, fmt.Errorf("invalid format")
	}
	p := blob[len(magic):]
	salt, rest := p[:saltLen], p[saltLen:]
	gcm, err := newGCM(pass, salt)
	if err != nil {
		return nil, err
	}
	ns := gcm.NonceSize()
	if len(rest) < ns {
		return nil, fmt.Errorf("invalid format")
	}
	nonce, ct := rest[:ns], rest[ns:]
	pt, err := gcm.Open(nil, nonce, ct, []byte(magic))
	if err != nil {
		return nil, fmt.Errorf("wrong password (or tampered file)")
	}
	return pt, nil
}

func main() {
	pass := os.Getenv("PASSWORD")
	if pass == "" {
		fatal(fmt.Errorf("set the PASSWORD environment variable"))
	}
	if len(os.Args) < 2 {
		fatal(fmt.Errorf("usage: %s [encrypt <text> | decrypt]", os.Args[0]))
	}

	switch os.Args[1] {
	case "encrypt":
		if len(os.Args) < 3 {
			fatal(fmt.Errorf("usage: %s encrypt <text>", os.Args[0]))
		}
		blob, err := encrypt(pass, []byte(os.Args[2]))
		if err != nil {
			fatal(err)
		}
		_, _ = os.Stdout.Write(blob)
	case "decrypt":
		blob, err := io.ReadAll(os.Stdin)
		if err != nil {
			fatal(err)
		}
		plain, err := decrypt(pass, blob)
		if err != nil {
			fatal(err)
		}
		fmt.Println(string(plain))
	default:
		fatal(fmt.Errorf("unknown command: %s", os.Args[1]))
	}
}