// VaultZip — password-encrypted archive with built-in checksum verification.
//
// Usage:
//
//	vaultzip pack   <path> [<path> ...] -o vault.vz --password PASS [--dry-run]
//	vaultzip unpack <vault.vz> -o outdir --password PASS
//	vaultzip verify <vault.vz> --password PASS
//
// "pack --dry-run" lists what would go into the vault, with sizes, and writes
// nothing — the read-only half of packing, so that a caller with something to
// lose can look before it commits.
//
// A vault is tar+gzip, then AES-256-GCM encrypted with a key stretched
// from the password via iterated SHA-256 (a hand-rolled, dependency-free
// stand-in for PBKDF2 — stdlib has no KDF — see the build plan / README
// before trusting this for anything beyond a prototype: it has not been
// audited). "verify" authenticates and lists contents with per-file
// SHA-256 without writing anything to disk.
package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

const (
	magic     = "VLTZ1\n"
	saltLen   = 16
	nonceLen  = 12
	kdfRounds = 200000
)

func main() {
	if len(os.Args) < 2 {
		// Double-clicked in Explorer rather than run from a prompt: ask what
		// the program needs and stay on screen. Printing usage and exiting
		// here is what made the window vanish instantly.
		if interactiveConsole() {
			runGuided()
			return
		}
		usage()
		os.Exit(1)
	}
	switch os.Args[1] {
	case "pack":
		cmdPack(os.Args[2:])
	case "unpack":
		cmdUnpack(os.Args[2:], true)
	case "verify":
		cmdUnpack(os.Args[2:], false)
	case "-h", "--help", "help":
		usage()
	default:
		fmt.Fprintf(os.Stderr, "unknown command %q\n\n", os.Args[1])
		usage()
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprint(os.Stderr, `VaultZip — password-encrypted archive with checksum verification

Usage:
  vaultzip pack   <path> [<path> ...] -o vault.vz --password PASS [--dry-run]
  vaultzip unpack <vault.vz> -o outdir --password PASS
  vaultzip verify <vault.vz> --password PASS

Flags:
  -o <file>     Vault to write (pack) or directory to extract into (unpack).
  --password P  Vault password. Required by every command.
  --dry-run     pack only. List every file that would go into the vault, with
                its size on disk, and write nothing at all — no vault, no
                partial file, no temporary. Warns if a vault already exists at
                the -o path and would be replaced. Add it to the command you
                were about to run to see what it would do first.

"verify" authenticates and lists contents without extracting anything.
`)
}

func reorderFlags(args []string, valueFlags map[string]bool) []string {
	var flags, positional []string
	for i := 0; i < len(args); i++ {
		a := args[i]
		name := strings.TrimLeft(a, "-")
		if strings.HasPrefix(a, "-") && valueFlags[name] {
			flags = append(flags, a)
			if i+1 < len(args) {
				i++
				flags = append(flags, args[i])
			}
			continue
		}
		if strings.HasPrefix(a, "-") {
			flags = append(flags, a)
			continue
		}
		positional = append(positional, a)
	}
	return append(flags, positional...)
}

func stretchKey(password string, salt []byte) []byte {
	key := append([]byte(password), salt...)
	sum := sha256.Sum256(key)
	for i := 0; i < kdfRounds; i++ {
		sum = sha256.Sum256(sum[:])
	}
	return sum[:]
}

// eachPackFile walks the pack targets in order and calls visit once per file
// that would be stored, with the name it would be stored under and an open
// handle to it.
//
// Packing and its dry run both go through here, and that is the whole point of
// the function existing: the list the dry run prints is produced by the same
// walk, in the same order, with the same relative names, and skips the same
// unreadable files for the same reasons. A dry run assembled by a second,
// similar-looking loop would eventually disagree with the real one, and a
// preview that disagrees with what follows it is worse than no preview.
func eachPackFile(targets []string, visit func(rel string, f *os.File, fi os.FileInfo) error) error {
	for _, t := range targets {
		if _, err := os.Stat(t); err != nil {
			fmt.Fprintf(os.Stderr, "skip %s: %v\n", t, err)
			continue
		}
		base := filepath.Dir(t)
		walkErr := filepath.Walk(t, func(path string, fi os.FileInfo, err error) error {
			if err != nil || fi.IsDir() {
				return nil
			}
			rel, err := filepath.Rel(base, path)
			if err != nil {
				rel = filepath.Base(path)
			}
			// Opened even by the dry run, so that a file the real pack would
			// have to skip is skipped in the preview too rather than being
			// promised and then quietly dropped.
			f, err := os.Open(path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "skip %s: %v\n", path, err)
				return nil
			}
			defer f.Close()
			return visit(filepath.ToSlash(rel), f, fi)
		})
		if walkErr != nil {
			return fmt.Errorf("error walking %s: %w", t, walkErr)
		}
	}
	return nil
}

// packDryRun lists what packing would store and writes nothing whatsoever — no
// vault, no partial file, no temporary. It stops before the key stretching as
// well as before the output file, so it is quick enough to sit in front of a
// real pack rather than doubling its cost.
//
// It reports sizes as they are on disk, before compression, and says so instead
// of guessing at a compressed figure it has not computed.
func packDryRun(targets []string, out string) error {
	var fileCount int
	var totalBytes int64
	err := eachPackFile(targets, func(rel string, _ *os.File, fi os.FileInfo) error {
		fileCount++
		totalBytes += fi.Size()
		fmt.Printf("  would add  %s (%s)\n", rel, humanBytes(fi.Size()))
		return nil
	})
	if err != nil {
		return err
	}
	if fileCount == 0 {
		fmt.Println("Nothing to pack.")
		return nil
	}
	fmt.Printf("\nWould write %s: %d files, %s before compression.\n", out, fileCount, humanBytes(totalBytes))
	if info, statErr := os.Stat(out); statErr == nil && !info.IsDir() {
		fmt.Printf("Warning: a vault already exists at %s (%s) and would be replaced.\n", out, humanBytes(info.Size()))
	}
	fmt.Println("Dry run: nothing has been written.")
	return nil
}

func cmdPack(args []string) {
	fs := flag.NewFlagSet("pack", flag.ExitOnError)
	out := fs.String("o", "", "output vault file (required)")
	password := fs.String("password", "", "vault password (required)")
	dryRun := fs.Bool("dry-run", false, "list what would go into the vault and write nothing")
	fs.Parse(reorderFlags(args, map[string]bool{"o": true, "password": true}))
	targets := fs.Args()
	if len(targets) == 0 || *out == "" || *password == "" {
		fmt.Fprintln(os.Stderr, "usage: vaultzip pack <path> [<path> ...] -o vault.vz --password PASS [--dry-run]")
		os.Exit(1)
	}

	// Deliberately before anything is created, and before the 200,000-round key
	// stretch: a dry run must be able to run on a folder the customer has no
	// intention of packing yet, and cost them nothing for looking.
	if *dryRun {
		if err := packDryRun(targets, *out); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}

	var tarBuf bytes.Buffer
	gz := gzip.NewWriter(&tarBuf)
	tw := tar.NewWriter(gz)

	var fileCount int
	var totalBytes int64
	if err := eachPackFile(targets, func(rel string, f *os.File, fi os.FileInfo) error {
		hdr := &tar.Header{Name: rel, Size: fi.Size(), Mode: int64(fi.Mode().Perm()), ModTime: fi.ModTime()}
		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		n, err := io.Copy(tw, f)
		if err != nil {
			return err
		}
		fileCount++
		totalBytes += n
		fmt.Printf("  add  %s (%s)\n", rel, humanBytes(n))
		return nil
	}); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if fileCount == 0 {
		fmt.Println("Nothing to pack.")
		return
	}
	tw.Close()
	gz.Close()

	salt := make([]byte, saltLen)
	rand.Read(salt)
	key := stretchKey(*password, salt)
	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, "cipher error:", err)
		os.Exit(1)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		fmt.Fprintln(os.Stderr, "gcm error:", err)
		os.Exit(1)
	}
	nonce := make([]byte, nonceLen)
	rand.Read(nonce)
	ciphertext := gcm.Seal(nil, nonce, tarBuf.Bytes(), nil)

	f, err := os.Create(*out)
	if err != nil {
		fmt.Fprintln(os.Stderr, "error creating vault:", err)
		os.Exit(1)
	}
	defer f.Close()
	f.WriteString(magic)
	f.Write(salt)
	f.Write(nonce)
	f.Write(ciphertext)

	fmt.Printf("\nVault written: %s (%d files, %s plaintext, %s on disk)\n", *out, fileCount, humanBytes(totalBytes), humanBytes(int64(len(ciphertext)+len(magic)+saltLen+nonceLen)))
}

func cmdUnpack(args []string, extract bool) {
	name := "unpack"
	if !extract {
		name = "verify"
	}
	fs := flag.NewFlagSet(name, flag.ExitOnError)
	out := fs.String("o", "", "output directory (required for unpack)")
	password := fs.String("password", "", "vault password (required)")
	fs.Parse(reorderFlags(args, map[string]bool{"o": true, "password": true}))
	pos := fs.Args()
	if len(pos) != 1 || *password == "" || (extract && *out == "") {
		fmt.Fprintf(os.Stderr, "usage: vaultzip %s <vault.vz> %s--password PASS\n", name, map[bool]string{true: "-o outdir ", false: ""}[extract])
		os.Exit(1)
	}
	vaultPath := pos[0]

	raw, err := os.ReadFile(vaultPath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "error reading vault:", err)
		os.Exit(1)
	}
	if len(raw) < len(magic)+saltLen+nonceLen || string(raw[:len(magic)]) != magic {
		fmt.Fprintln(os.Stderr, "not a VaultZip file (bad magic header)")
		os.Exit(1)
	}
	off := len(magic)
	salt := raw[off : off+saltLen]
	off += saltLen
	nonce := raw[off : off+nonceLen]
	off += nonceLen
	ciphertext := raw[off:]

	key := stretchKey(*password, salt)
	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Fprintln(os.Stderr, "cipher error:", err)
		os.Exit(1)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		fmt.Fprintln(os.Stderr, "gcm error:", err)
		os.Exit(1)
	}
	plain, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, "AUTHENTICATION FAILED — wrong password or the vault has been tampered with")
		os.Exit(2)
	}

	gz, err := gzip.NewReader(bytes.NewReader(plain))
	if err != nil {
		fmt.Fprintln(os.Stderr, "corrupt archive stream:", err)
		os.Exit(1)
	}
	tr := tar.NewReader(gz)

	var fileCount int
	var totalBytes int64
	action := "verified"
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			fmt.Fprintln(os.Stderr, "corrupt archive entry:", err)
			os.Exit(1)
		}
		h := sha256.New()
		var w io.Writer = h
		var destFile *os.File
		if extract {
			dest := filepath.Join(*out, filepath.FromSlash(hdr.Name))
			if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil {
				fmt.Fprintln(os.Stderr, "error creating dir:", err)
				os.Exit(1)
			}
			destFile, err = os.OpenFile(dest, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.FileMode(hdr.Mode))
			if err != nil {
				fmt.Fprintln(os.Stderr, "error creating file:", err)
				os.Exit(1)
			}
			w = io.MultiWriter(h, destFile)
			action = "extracted"
		}
		n, err := io.Copy(w, tr)
		if destFile != nil {
			destFile.Close()
		}
		if err != nil {
			fmt.Fprintln(os.Stderr, "error reading entry:", err)
			os.Exit(1)
		}
		fileCount++
		totalBytes += n
		fmt.Printf("  %-10s %8s  sha256:%s  %s\n", action, humanBytes(n), hex.EncodeToString(h.Sum(nil))[:16], hdr.Name)
	}

	verb := "Verified"
	if extract {
		verb = "Extracted"
	}
	fmt.Printf("\n%s: password correct, integrity confirmed, %d files, %s\n", verb, fileCount, humanBytes(totalBytes))
}

func humanBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for x := n / unit; x >= unit; x /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(n)/float64(div), "KMGTPE"[exp])
}