// Command partitionguard is a read-only MBR and GPT partition table parser.
//
// It decodes the raw bytes of a disk image and reports the partition layout,
// validates GPT CRC32 checksums, and performs structural sanity checks.
// It never writes to the image it inspects.
package main

import (
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"hash/crc32"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"unicode/utf16"
)

const (
	gptSignature   = "EFI PART"
	mbrPartOffset  = 446
	mbrEntrySize   = 16
	mbrSigOffset   = 510
	gptMinHeader   = 92
	maxArrayBytes  = 16 << 20
	exitUsage      = 1
	exitCheckFails = 2
)

// ---------------------------------------------------------------------------
// Shared Techlosoft UX helpers.
// ---------------------------------------------------------------------------

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 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])
}

// ---------------------------------------------------------------------------
// Image access. Every read is bounds-checked; the file is opened read-only.
// ---------------------------------------------------------------------------

type image struct {
	path       string
	f          *os.File
	size       int64
	sectorSize int
}

func openImage(path string, sectorSize int) (*image, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	st, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, err
	}
	if st.IsDir() {
		f.Close()
		return nil, fmt.Errorf("%s is a directory, not a disk image", path)
	}
	if st.Size() == 0 {
		f.Close()
		return nil, fmt.Errorf("%s is empty (0 bytes): not a disk image", path)
	}
	if st.Size() < int64(sectorSize) {
		f.Close()
		return nil, fmt.Errorf("%s is only %d bytes, smaller than one %d-byte sector: too small to hold a partition table",
			path, st.Size(), sectorSize)
	}
	return &image{path: path, f: f, size: st.Size(), sectorSize: sectorSize}, nil
}

func (im *image) Close() error { return im.f.Close() }

func (im *image) sectors() uint64 { return uint64(im.size) / uint64(im.sectorSize) }

func (im *image) readAt(off int64, n int) ([]byte, error) {
	if off < 0 || n < 0 {
		return nil, fmt.Errorf("invalid read (offset %d, length %d)", off, n)
	}
	if off > im.size || off+int64(n) > im.size {
		return nil, fmt.Errorf("read of %d bytes at offset %d extends past end of image (%d bytes)", n, off, im.size)
	}
	buf := make([]byte, n)
	if _, err := io.ReadFull(io.NewSectionReader(im.f, off, int64(n)), buf); err != nil {
		return nil, fmt.Errorf("read at offset %d: %w", off, err)
	}
	return buf, nil
}

// lbaOffset converts an LBA to a byte offset, guarding against overflow.
func (im *image) lbaOffset(lba uint64) (int64, bool) {
	const maxOff = uint64(1) << 62
	if uint64(im.sectorSize) == 0 || lba > maxOff/uint64(im.sectorSize) {
		return 0, false
	}
	return int64(lba) * int64(im.sectorSize), true
}

func (im *image) readLBA(lba uint64, n int) ([]byte, error) {
	off, ok := im.lbaOffset(lba)
	if !ok {
		return nil, fmt.Errorf("LBA %d is out of representable range", lba)
	}
	return im.readAt(off, n)
}

// ---------------------------------------------------------------------------
// Little-endian scalar readers.
// ---------------------------------------------------------------------------

func le16(b []byte) uint16 { return binary.LittleEndian.Uint16(b) }
func le32(b []byte) uint32 { return binary.LittleEndian.Uint32(b) }
func le64(b []byte) uint64 { return binary.LittleEndian.Uint64(b) }

// formatGUID renders a 16-byte mixed-endian GUID the way Microsoft and the
// UEFI spec do: the first three groups are little-endian, the last two are
// byte-ordered as stored.
func formatGUID(b []byte) string {
	if len(b) < 16 {
		return ""
	}
	return strings.ToUpper(fmt.Sprintf("%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x",
		le32(b[0:4]), le16(b[4:6]), le16(b[6:8]),
		b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15]))
}

func isZeroGUID(b []byte) bool {
	for _, c := range b[:16] {
		if c != 0 {
			return false
		}
	}
	return true
}

// decodeUTF16LEName decodes a NUL-padded UTF-16LE partition name.
func decodeUTF16LEName(b []byte) string {
	u := make([]uint16, 0, len(b)/2)
	for i := 0; i+1 < len(b); i += 2 {
		c := le16(b[i : i+2])
		if c == 0 {
			break
		}
		u = append(u, c)
	}
	return string(utf16.Decode(u))
}

// ---------------------------------------------------------------------------
// Type tables.
// ---------------------------------------------------------------------------

var mbrTypeNames = map[byte]string{
	0x00: "Empty",
	0x01: "FAT12",
	0x04: "FAT16 <32M",
	0x05: "Extended",
	0x06: "FAT16",
	0x07: "NTFS/exFAT",
	0x0B: "FAT32 (CHS)",
	0x0C: "FAT32 (LBA)",
	0x0E: "FAT16 (LBA)",
	0x0F: "Extended (LBA)",
	0x82: "Linux swap",
	0x83: "Linux",
	0x8E: "Linux LVM",
	0xA5: "FreeBSD",
	0xEE: "GPT protective",
	0xEF: "EFI System",
	0xFD: "Linux RAID autodetect",
}

func mbrTypeName(t byte) string {
	if n, ok := mbrTypeNames[t]; ok {
		return n
	}
	return "Unknown"
}

var gptTypeNames = map[string]string{
	"C12A7328-F81F-11D2-BA4B-00A0C93EC93B": "EFI System",
	"EBD0A0A2-B9E5-4433-87C0-68B6B72699C7": "Microsoft Basic Data",
	"E3C9E316-0B5C-4DB8-817D-F92DF00215AE": "Microsoft Reserved",
	"0FC63DAF-8483-4772-8E79-3D69D8477DE4": "Linux filesystem",
	"0657FD6D-A4AB-43C4-84E5-0933C84B4F4F": "Linux swap",
	"21686148-6449-6E6F-744E-656564454649": "BIOS boot",
	"E6D6D379-F507-44C2-A23C-238F2A3DF928": "Linux LVM",
	"A19D880F-05FC-4D3B-A006-743F0F84911E": "Linux RAID",
	"933AC7E1-2EB4-4F13-B844-0E14E2AEF915": "Linux /home",
	"DE94BBA4-06D1-4D40-A16A-BFD50179D6AC": "Windows Recovery",
	"48465300-0000-11AA-AA11-00306543ECAC": "Apple HFS+",
	"7C3457EF-0000-11AA-AA11-00306543ECAC": "Apple APFS",
}

func gptTypeName(guid string) string {
	if n, ok := gptTypeNames[guid]; ok {
		return n
	}
	return "Unknown"
}

var gptAttrNames = []struct {
	bit  uint
	name string
}{
	{0, "RequiredPartition"},
	{1, "NoBlockIOProtocol"},
	{2, "LegacyBIOSBootable"},
	{60, "ReadOnly"},
	{61, "ShadowCopy"},
	{62, "Hidden"},
	{63, "NoAutomount"},
}

func attrFlagNames(a uint64) []string {
	var out []string
	for _, f := range gptAttrNames {
		if a&(uint64(1)<<f.bit) != 0 {
			out = append(out, f.name)
		}
	}
	return out
}

// ---------------------------------------------------------------------------
// MBR parsing.
// ---------------------------------------------------------------------------

type mbrPartition struct {
	Index    int    `json:"index"`
	Bootable bool   `json:"bootable"`
	Type     byte   `json:"-"`
	TypeCode string `json:"type_code"`
	TypeName string `json:"type_name"`
	StartLBA uint32 `json:"start_lba"`
	EndLBA   uint64 `json:"end_lba"`
	Sectors  uint32 `json:"sectors"`
	Bytes    int64  `json:"size_bytes"`
	Size     string `json:"size_human"`
}

type mbrInfo struct {
	SignaturePresent bool           `json:"signature_present"`
	Signature        string         `json:"signature"`
	Protective       bool           `json:"protective"`
	Partitions       []mbrPartition `json:"partitions"`
}

func parseMBR(im *image) (*mbrInfo, error) {
	sec, err := im.readAt(0, 512)
	if err != nil {
		return nil, fmt.Errorf("cannot read boot sector: %w", err)
	}
	mi := &mbrInfo{
		Signature: fmt.Sprintf("0x%02X%02X", sec[mbrSigOffset], sec[mbrSigOffset+1]),
	}
	mi.SignaturePresent = sec[mbrSigOffset] == 0x55 && sec[mbrSigOffset+1] == 0xAA
	for i := 0; i < 4; i++ {
		e := sec[mbrPartOffset+i*mbrEntrySize : mbrPartOffset+(i+1)*mbrEntrySize]
		typ := e[4]
		start := le32(e[8:12])
		count := le32(e[12:16])
		if typ == 0xEE {
			mi.Protective = true
		}
		if typ == 0x00 || count == 0 {
			continue
		}
		p := mbrPartition{
			Index:    i + 1,
			Bootable: e[0] == 0x80,
			Type:     typ,
			TypeCode: fmt.Sprintf("0x%02X", typ),
			TypeName: mbrTypeName(typ),
			StartLBA: start,
			Sectors:  count,
			EndLBA:   uint64(start) + uint64(count) - 1,
			Bytes:    int64(count) * int64(im.sectorSize),
		}
		p.Size = humanBytes(p.Bytes)
		mi.Partitions = append(mi.Partitions, p)
	}
	return mi, nil
}

// ---------------------------------------------------------------------------
// GPT parsing.
// ---------------------------------------------------------------------------

type crcCheck struct {
	Stored   string `json:"stored"`
	Computed string `json:"computed"`
	Status   string `json:"status"`
	ok       bool
}

func mkCRC(stored, computed uint32) crcCheck {
	c := crcCheck{
		Stored:   fmt.Sprintf("0x%08X", stored),
		Computed: fmt.Sprintf("0x%08X", computed),
		ok:       stored == computed,
	}
	if c.ok {
		c.Status = "VALID"
	} else {
		c.Status = "MISMATCH"
	}
	return c
}

type gptPartition struct {
	Index      int      `json:"index"`
	TypeGUID   string   `json:"type_guid"`
	TypeName   string   `json:"type_name"`
	PartGUID   string   `json:"partition_guid"`
	FirstLBA   uint64   `json:"start_lba"`
	LastLBA    uint64   `json:"end_lba"`
	Sectors    uint64   `json:"sectors"`
	Bytes      int64    `json:"size_bytes"`
	Size       string   `json:"size_human"`
	Attributes string   `json:"attributes"`
	AttrFlags  []string `json:"attribute_flags"`
	Name       string   `json:"name"`
}

type gptHeader struct {
	Revision       string `json:"revision"`
	HeaderSize     uint32 `json:"header_size"`
	CurrentLBA     uint64 `json:"current_lba"`
	BackupLBA      uint64 `json:"backup_lba"`
	FirstUsableLBA uint64 `json:"first_usable_lba"`
	LastUsableLBA  uint64 `json:"last_usable_lba"`
	DiskGUID       string `json:"disk_guid"`
	PartEntryLBA   uint64 `json:"partition_entry_lba"`
	NumEntries     uint32 `json:"entry_count"`
	EntrySize      uint32 `json:"entry_size"`

	storedHeaderCRC uint32
	storedArrayCRC  uint32
	sane            bool
}

type gptInfo struct {
	Header     *gptHeader     `json:"header"`
	HeaderCRC  crcCheck       `json:"header_crc32"`
	ArrayCRC   *crcCheck      `json:"array_crc32"`
	ArrayError string         `json:"array_error,omitempty"`
	Partitions []gptPartition `json:"partitions"`
}

// parseGPTHeaderSector decodes a GPT header out of a sector-sized buffer.
// It returns (nil, nil) when the signature is absent — that is not an error,
// it simply means there is no GPT here.
func parseGPTHeaderSector(sec []byte) (*gptHeader, crcCheck) {
	var none crcCheck
	if len(sec) < gptMinHeader || string(sec[0:8]) != gptSignature {
		return nil, none
	}
	h := &gptHeader{
		Revision:        fmt.Sprintf("%d.%d", le16(sec[10:12]), le16(sec[8:10])),
		HeaderSize:      le32(sec[12:16]),
		CurrentLBA:      le64(sec[24:32]),
		BackupLBA:       le64(sec[32:40]),
		FirstUsableLBA:  le64(sec[40:48]),
		LastUsableLBA:   le64(sec[48:56]),
		DiskGUID:        formatGUID(sec[56:72]),
		PartEntryLBA:    le64(sec[72:80]),
		NumEntries:      le32(sec[80:84]),
		EntrySize:       le32(sec[84:88]),
		storedHeaderCRC: le32(sec[16:20]),
		storedArrayCRC:  le32(sec[88:92]),
	}
	h.sane = h.HeaderSize >= gptMinHeader && int(h.HeaderSize) <= len(sec)
	if !h.sane {
		// Cannot recompute a CRC over a header of implausible length.
		return h, crcCheck{
			Stored:   fmt.Sprintf("0x%08X", h.storedHeaderCRC),
			Computed: "n/a",
			Status:   "MISMATCH",
		}
	}
	tmp := make([]byte, h.HeaderSize)
	copy(tmp, sec[:h.HeaderSize])
	// The header CRC field itself must be zeroed before recomputing.
	for i := 16; i < 20; i++ {
		tmp[i] = 0
	}
	return h, mkCRC(h.storedHeaderCRC, crc32.ChecksumIEEE(tmp))
}

// readGPTArray reads and decodes the partition entry array.
func readGPTArray(im *image, h *gptHeader) (crcCheck, []gptPartition, error) {
	if h.EntrySize < 128 || h.EntrySize > 32768 || h.EntrySize%128 != 0 {
		return crcCheck{}, nil, fmt.Errorf("implausible partition entry size %d (must be a multiple of 128 between 128 and 32768)", h.EntrySize)
	}
	if h.NumEntries == 0 {
		return crcCheck{}, nil, fmt.Errorf("header declares 0 partition entries")
	}
	total := uint64(h.NumEntries) * uint64(h.EntrySize)
	if total > maxArrayBytes {
		return crcCheck{}, nil, fmt.Errorf("partition array of %d entries x %d bytes = %d bytes is implausibly large", h.NumEntries, h.EntrySize, total)
	}
	off, ok := im.lbaOffset(h.PartEntryLBA)
	if !ok {
		return crcCheck{}, nil, fmt.Errorf("partition entry LBA %d is out of range", h.PartEntryLBA)
	}
	buf, err := im.readAt(off, int(total))
	if err != nil {
		return crcCheck{}, nil, fmt.Errorf("cannot read partition array at LBA %d: %w", h.PartEntryLBA, err)
	}
	cc := mkCRC(h.storedArrayCRC, crc32.ChecksumIEEE(buf))

	var parts []gptPartition
	for i := 0; i < int(h.NumEntries); i++ {
		e := buf[i*int(h.EntrySize) : (i+1)*int(h.EntrySize)]
		if isZeroGUID(e[0:16]) {
			continue // unused entry
		}
		first := le64(e[32:40])
		last := le64(e[40:48])
		attr := le64(e[48:56])
		nameEnd := 128
		if len(e) < nameEnd {
			nameEnd = len(e)
		}
		guid := formatGUID(e[0:16])
		p := gptPartition{
			Index:      i + 1,
			TypeGUID:   guid,
			TypeName:   gptTypeName(guid),
			PartGUID:   formatGUID(e[16:32]),
			FirstLBA:   first,
			LastLBA:    last,
			Attributes: fmt.Sprintf("0x%016X", attr),
			AttrFlags:  attrFlagNames(attr),
			Name:       decodeUTF16LEName(e[56:nameEnd]),
		}
		if last >= first {
			p.Sectors = last - first + 1
			if p.Sectors <= uint64(1)<<40 {
				p.Bytes = int64(p.Sectors) * int64(im.sectorSize)
				p.Size = humanBytes(p.Bytes)
			} else {
				p.Size = "out of range"
			}
		} else {
			p.Size = "invalid (end < start)"
		}
		parts = append(parts, p)
	}
	return cc, parts, nil
}

// loadGPT returns nil (with no error) when the image simply has no GPT.
func loadGPT(im *image) *gptInfo {
	sec, err := im.readLBA(1, im.sectorSize)
	if err != nil {
		// The image is too small to hold a header at LBA 1: no GPT.
		return nil
	}
	h, hcrc := parseGPTHeaderSector(sec)
	if h == nil {
		return nil
	}
	gi := &gptInfo{Header: h, HeaderCRC: hcrc}
	acrc, parts, aerr := readGPTArray(im, h)
	if aerr != nil {
		gi.ArrayError = aerr.Error()
	} else {
		gi.ArrayCRC = &acrc
		gi.Partitions = parts
	}
	return gi
}

// ---------------------------------------------------------------------------
// Scheme detection.
// ---------------------------------------------------------------------------

type layout struct {
	im     *image
	mbr    *mbrInfo
	gpt    *gptInfo
	scheme string
	desc   string
}

func analyze(im *image) (*layout, error) {
	mbr, err := parseMBR(im)
	if err != nil {
		return nil, err
	}
	gpt := loadGPT(im)
	l := &layout{im: im, mbr: mbr, gpt: gpt}
	switch {
	case gpt != nil && mbr.Protective:
		l.scheme = "gpt"
		l.desc = "GPT (with protective MBR)"
	case gpt != nil && mbr.SignaturePresent:
		l.scheme = "gpt"
		l.desc = "GPT (MBR signature present but no 0xEE protective entry)"
	case gpt != nil:
		l.scheme = "gpt"
		l.desc = "GPT (no protective MBR)"
	case mbr.SignaturePresent && mbr.Protective:
		l.scheme = "mbr"
		l.desc = "MBR carrying a 0xEE protective entry, but no readable GPT header at LBA 1 (GPT missing or damaged)"
	case mbr.SignaturePresent && len(mbr.Partitions) > 0:
		l.scheme = "mbr"
		l.desc = "MBR"
	case mbr.SignaturePresent:
		l.scheme = "mbr-empty"
		l.desc = "MBR (valid signature, no partitions defined)"
	default:
		l.scheme = "none"
		l.desc = "none (no MBR signature and no GPT header)"
	}
	return l, nil
}

// ---------------------------------------------------------------------------
// info
// ---------------------------------------------------------------------------

type infoJSON struct {
	Image      string   `json:"image"`
	SizeBytes  int64    `json:"size_bytes"`
	SizeHuman  string   `json:"size_human"`
	SectorSize int      `json:"sector_size"`
	Sectors    uint64   `json:"sectors"`
	Scheme     string   `json:"scheme"`
	SchemeDesc string   `json:"scheme_description"`
	MBR        *mbrInfo `json:"mbr"`
	GPT        *gptInfo `json:"gpt"`
	Error      string   `json:"error,omitempty"`
}

// arrayErr reports the reason the GPT partition array could not be read, if any.
func (l *layout) arrayErr() string {
	if l.gpt != nil {
		return l.gpt.ArrayError
	}
	return ""
}

func cmdInfo(l *layout, asJSON bool) int {
	im := l.im
	if asJSON {
		out := infoJSON{
			Image:      im.path,
			SizeBytes:  im.size,
			SizeHuman:  humanBytes(im.size),
			SectorSize: im.sectorSize,
			Sectors:    im.sectors(),
			Scheme:     l.scheme,
			SchemeDesc: l.desc,
			MBR:        l.mbr,
			GPT:        l.gpt,
			Error:      l.arrayErr(),
		}
		if rc := emitJSON(out); rc != 0 {
			return rc
		}
		if out.Error != "" {
			return exitUsage
		}
		return 0
	}

	fmt.Printf("Image:        %s\n", im.path)
	fmt.Printf("Size:         %s (%d bytes)\n", humanBytes(im.size), im.size)
	fmt.Printf("Sector size:  %d bytes\n", im.sectorSize)
	fmt.Printf("Sectors:      %d\n", im.sectors())
	fmt.Printf("Scheme:       %s\n", l.desc)
	fmt.Println()
	fmt.Println("MBR:")
	fmt.Printf("  Boot signature:   %s (%s)\n", l.mbr.Signature, yesNo(l.mbr.SignaturePresent, "present", "ABSENT"))
	fmt.Printf("  Protective 0xEE:  %s\n", yesNo(l.mbr.Protective, "yes", "no"))
	fmt.Printf("  Entries in use:   %d\n", len(l.mbr.Partitions))

	if l.gpt == nil {
		fmt.Println()
		fmt.Println("GPT:  not present (no \"EFI PART\" signature at LBA 1)")
		return 0
	}
	h := l.gpt.Header
	fmt.Println()
	fmt.Println("GPT header (LBA 1):")
	fmt.Printf("  Revision:             %s\n", h.Revision)
	fmt.Printf("  Header size:          %d bytes\n", h.HeaderSize)
	fmt.Printf("  Current LBA:          %d\n", h.CurrentLBA)
	fmt.Printf("  Backup LBA:           %d\n", h.BackupLBA)
	fmt.Printf("  First usable LBA:     %d\n", h.FirstUsableLBA)
	fmt.Printf("  Last usable LBA:      %d\n", h.LastUsableLBA)
	fmt.Printf("  Disk GUID:            %s\n", h.DiskGUID)
	fmt.Printf("  Partition entry LBA:  %d\n", h.PartEntryLBA)
	fmt.Printf("  Entry count:          %d\n", h.NumEntries)
	fmt.Printf("  Entry size:           %d bytes\n", h.EntrySize)
	fmt.Printf("  Header CRC32:         stored %s  computed %s  %s\n",
		l.gpt.HeaderCRC.Stored, l.gpt.HeaderCRC.Computed, l.gpt.HeaderCRC.Status)
	if l.gpt.ArrayCRC != nil {
		fmt.Printf("  Array CRC32:          stored %s  computed %s  %s\n",
			l.gpt.ArrayCRC.Stored, l.gpt.ArrayCRC.Computed, l.gpt.ArrayCRC.Status)
	} else {
		fmt.Printf("  Array CRC32:          unreadable: %s\n", l.gpt.ArrayError)
	}
	fmt.Printf("  Partitions in use:    %d\n", len(l.gpt.Partitions))
	if l.gpt.ArrayError != "" {
		fmt.Fprintf(os.Stderr, "partitionguard: %s: partition array unreadable: %s\n", im.path, l.gpt.ArrayError)
		return exitUsage
	}
	return 0
}

func yesNo(b bool, y, n string) string {
	if b {
		return y
	}
	return n
}

// ---------------------------------------------------------------------------
// list
// ---------------------------------------------------------------------------

type listJSON struct {
	Image         string         `json:"image"`
	Scheme        string         `json:"scheme"`
	SchemeDesc    string         `json:"scheme_description"`
	SectorSize    int            `json:"sector_size"`
	MBRPartitions []mbrPartition `json:"mbr_partitions,omitempty"`
	GPTPartitions []gptPartition `json:"gpt_partitions,omitempty"`
	Error         string         `json:"error,omitempty"`
}

func cmdList(l *layout, asJSON bool) int {
	out := listJSON{
		Image:      l.im.path,
		Scheme:     l.scheme,
		SchemeDesc: l.desc,
		SectorSize: l.im.sectorSize,
		Error:      l.arrayErr(),
	}
	if l.gpt != nil {
		out.GPTPartitions = l.gpt.Partitions
	} else {
		out.MBRPartitions = l.mbr.Partitions
	}
	if asJSON {
		if out.GPTPartitions == nil {
			out.GPTPartitions = []gptPartition{}
		}
		if out.MBRPartitions == nil {
			out.MBRPartitions = []mbrPartition{}
		}
		if rc := emitJSON(out); rc != 0 {
			return rc
		}
		if out.Error != "" {
			return exitUsage
		}
		return 0
	}

	fmt.Printf("Image:  %s\n", l.im.path)
	fmt.Printf("Scheme: %s\n", l.desc)
	fmt.Printf("Sector: %d bytes\n", l.im.sectorSize)
	fmt.Println()

	if l.gpt != nil {
		if l.gpt.ArrayError != "" {
			fmt.Fprintf(os.Stderr, "partitionguard: %s: cannot list partitions: %s\n", l.im.path, l.gpt.ArrayError)
			return exitUsage
		}
		if len(l.gpt.Partitions) == 0 {
			fmt.Println("No partitions defined.")
			return 0
		}
		fmt.Printf("%-3s %-12s %-12s %-12s %-12s %s\n", "#", "Start LBA", "End LBA", "Sectors", "Size", "Type")
		for _, p := range l.gpt.Partitions {
			fmt.Printf("%-3d %-12d %-12d %-12d %-12s %s\n",
				p.Index, p.FirstLBA, p.LastLBA, p.Sectors, p.Size, p.TypeName)
			fmt.Printf("    Type GUID:  %s\n", p.TypeGUID)
			fmt.Printf("    Part GUID:  %s\n", p.PartGUID)
			attrs := p.Attributes
			if len(p.AttrFlags) > 0 {
				attrs += " (" + strings.Join(p.AttrFlags, ", ") + ")"
			}
			fmt.Printf("    Attributes: %s\n", attrs)
			fmt.Printf("    Name:       %q\n", p.Name)
			fmt.Printf("    Bytes:      %d\n", p.Bytes)
			fmt.Println()
		}
		return 0
	}

	if len(l.mbr.Partitions) == 0 {
		fmt.Println("No partitions defined.")
		return 0
	}
	fmt.Printf("%-3s %-5s %-12s %-12s %-12s %-12s %-6s %s\n",
		"#", "Boot", "Start LBA", "End LBA", "Sectors", "Size", "Code", "Type")
	for _, p := range l.mbr.Partitions {
		fmt.Printf("%-3d %-5s %-12d %-12d %-12d %-12s %-6s %s\n",
			p.Index, yesNo(p.Bootable, "*", ""), p.StartLBA, p.EndLBA, p.Sectors, p.Size, p.TypeCode, p.TypeName)
	}
	return 0
}

// ---------------------------------------------------------------------------
// verify
// ---------------------------------------------------------------------------

type check struct {
	Name   string `json:"check"`
	Status string `json:"status"`
	Detail string `json:"detail"`
}

type verifyJSON struct {
	Image  string  `json:"image"`
	Scheme string  `json:"scheme"`
	Checks []check `json:"checks"`
	Passed int     `json:"passed"`
	Failed int     `json:"failed"`
	Skiped int     `json:"skipped"`
	Result string  `json:"result"`
}

type span struct {
	label string
	start uint64
	end   uint64
}

func cmdVerify(l *layout, asJSON bool) int {
	im := l.im
	var checks []check
	add := func(name, status, detail string) {
		checks = append(checks, check{Name: name, Status: status, Detail: detail})
	}

	// 1. MBR boot signature.
	if l.mbr.SignaturePresent {
		add("mbr-signature", "PASS", fmt.Sprintf("boot signature %s found at offset 510", l.mbr.Signature))
	} else {
		add("mbr-signature", "FAIL", fmt.Sprintf("expected 0x55AA at offset 510, found %s", l.mbr.Signature))
	}

	// 2/3. GPT header and array CRC32.
	if l.gpt == nil {
		add("gpt-header-crc32", "SKIP", "no GPT header present at LBA 1")
		add("gpt-array-crc32", "SKIP", "no GPT header present at LBA 1")
		add("gpt-backup-header", "SKIP", "no GPT header present at LBA 1")
	} else {
		h := l.gpt.Header
		if l.gpt.HeaderCRC.ok {
			add("gpt-header-crc32", "PASS", fmt.Sprintf("stored %s matches computed %s", l.gpt.HeaderCRC.Stored, l.gpt.HeaderCRC.Computed))
		} else {
			add("gpt-header-crc32", "FAIL", fmt.Sprintf("primary GPT header CRC32 mismatch: stored %s, computed %s", l.gpt.HeaderCRC.Stored, l.gpt.HeaderCRC.Computed))
		}
		switch {
		case l.gpt.ArrayCRC == nil:
			add("gpt-array-crc32", "FAIL", "partition array unreadable: "+l.gpt.ArrayError)
		case l.gpt.ArrayCRC.ok:
			add("gpt-array-crc32", "PASS", fmt.Sprintf("stored %s matches computed %s over %d entries x %d bytes",
				l.gpt.ArrayCRC.Stored, l.gpt.ArrayCRC.Computed, h.NumEntries, h.EntrySize))
		default:
			add("gpt-array-crc32", "FAIL", fmt.Sprintf("partition entry array CRC32 mismatch: header stores %s, array actually hashes to %s (the partition array has been altered)",
				l.gpt.ArrayCRC.Stored, l.gpt.ArrayCRC.Computed))
		}
		st, detail := checkBackup(im, h)
		add("gpt-backup-header", st, detail)
	}

	// 4. Overlaps.
	var spans []span
	if l.gpt != nil {
		for _, p := range l.gpt.Partitions {
			if p.LastLBA >= p.FirstLBA {
				label := fmt.Sprintf("#%d", p.Index)
				if p.Name != "" {
					label += fmt.Sprintf(" (%q)", p.Name)
				}
				spans = append(spans, span{label, p.FirstLBA, p.LastLBA})
			}
		}
	} else {
		for _, p := range l.mbr.Partitions {
			spans = append(spans, span{fmt.Sprintf("#%d", p.Index), uint64(p.StartLBA), p.EndLBA})
		}
	}
	if len(spans) == 0 {
		add("partition-overlap", "SKIP", "no partitions to compare")
		add("partition-bounds", "SKIP", "no partitions to compare")
	} else {
		sorted := make([]span, len(spans))
		copy(sorted, spans)
		sort.Slice(sorted, func(i, j int) bool { return sorted[i].start < sorted[j].start })
		var overlaps []string
		for i := 0; i+1 < len(sorted); i++ {
			for j := i + 1; j < len(sorted); j++ {
				if sorted[j].start > sorted[i].end {
					break
				}
				overlaps = append(overlaps, fmt.Sprintf("%s [%d-%d] overlaps %s [%d-%d]",
					sorted[i].label, sorted[i].start, sorted[i].end,
					sorted[j].label, sorted[j].start, sorted[j].end))
			}
		}
		if len(overlaps) == 0 {
			add("partition-overlap", "PASS", fmt.Sprintf("%d partitions, no overlapping LBA ranges", len(spans)))
		} else {
			add("partition-overlap", "FAIL", strings.Join(overlaps, "; "))
		}

		// 5. Bounds.
		total := im.sectors()
		var oob []string
		for _, s := range spans {
			if s.end >= total {
				oob = append(oob, fmt.Sprintf("%s ends at LBA %d but the image holds only %d sectors (last LBA %d)",
					s.label, s.end, total, total-1))
			}
		}
		if len(oob) == 0 {
			add("partition-bounds", "PASS", fmt.Sprintf("all %d partitions fit within %d sectors", len(spans), total))
		} else {
			add("partition-bounds", "FAIL", strings.Join(oob, "; "))
		}
	}

	passed, failed, skipped := 0, 0, 0
	for _, c := range checks {
		switch c.Status {
		case "PASS":
			passed++
		case "FAIL":
			failed++
		default:
			skipped++
		}
	}
	result := "PASS"
	if failed > 0 {
		result = "FAIL"
	}

	if asJSON {
		emitJSON(verifyJSON{
			Image:  im.path,
			Scheme: l.scheme,
			Checks: checks,
			Passed: passed,
			Failed: failed,
			Skiped: skipped,
			Result: result,
		})
	} else {
		fmt.Printf("Image:  %s\n", im.path)
		fmt.Printf("Scheme: %s\n", l.desc)
		fmt.Println()
		for _, c := range checks {
			fmt.Printf("[%-4s] %-20s %s\n", c.Status, c.Name, c.Detail)
		}
		fmt.Println()
		fmt.Printf("Result: %s  (%d passed, %d failed, %d skipped)\n", result, passed, failed, skipped)
	}
	if failed > 0 {
		return exitCheckFails
	}
	return 0
}

// checkBackup compares the primary GPT header against the backup at BackupLBA.
func checkBackup(im *image, h *gptHeader) (string, string) {
	if h.BackupLBA == 0 {
		return "FAIL", "primary header declares backup LBA 0, which cannot hold a backup GPT header"
	}
	last := im.sectors()
	if last == 0 || h.BackupLBA > last-1 {
		return "FAIL", fmt.Sprintf("backup GPT header should live at LBA %d but the image holds only %d sectors (last LBA %d)",
			h.BackupLBA, last, last-1)
	}
	sec, err := im.readLBA(h.BackupLBA, im.sectorSize)
	if err != nil {
		return "FAIL", fmt.Sprintf("cannot read backup GPT header at LBA %d: %v", h.BackupLBA, err)
	}
	b, bcrc := parseGPTHeaderSector(sec)
	if b == nil {
		return "FAIL", fmt.Sprintf("no %q signature at backup LBA %d (backup GPT header is missing or destroyed)", gptSignature, h.BackupLBA)
	}
	var probs []string
	if !bcrc.ok {
		probs = append(probs, fmt.Sprintf("backup header CRC32 mismatch: stored %s, computed %s", bcrc.Stored, bcrc.Computed))
	}
	if b.DiskGUID != h.DiskGUID {
		probs = append(probs, fmt.Sprintf("disk GUID differs (primary %s, backup %s)", h.DiskGUID, b.DiskGUID))
	}
	if b.CurrentLBA != h.BackupLBA {
		probs = append(probs, fmt.Sprintf("backup header's current LBA is %d, expected %d", b.CurrentLBA, h.BackupLBA))
	}
	if b.BackupLBA != h.CurrentLBA {
		probs = append(probs, fmt.Sprintf("backup header points to primary LBA %d, expected %d", b.BackupLBA, h.CurrentLBA))
	}
	if b.FirstUsableLBA != h.FirstUsableLBA || b.LastUsableLBA != h.LastUsableLBA {
		probs = append(probs, fmt.Sprintf("usable range differs (primary %d-%d, backup %d-%d)",
			h.FirstUsableLBA, h.LastUsableLBA, b.FirstUsableLBA, b.LastUsableLBA))
	}
	if b.NumEntries != h.NumEntries || b.EntrySize != h.EntrySize {
		probs = append(probs, fmt.Sprintf("entry geometry differs (primary %dx%d, backup %dx%d)",
			h.NumEntries, h.EntrySize, b.NumEntries, b.EntrySize))
	}
	if b.storedArrayCRC != h.storedArrayCRC {
		probs = append(probs, fmt.Sprintf("partition array CRC32 differs (primary 0x%08X, backup 0x%08X)",
			h.storedArrayCRC, b.storedArrayCRC))
	}
	if len(probs) > 0 {
		return "FAIL", "primary and backup GPT headers disagree: " + strings.Join(probs, "; ")
	}
	return "PASS", fmt.Sprintf("backup header at LBA %d agrees with primary (CRC32 %s, disk GUID %s)",
		h.BackupLBA, bcrc.Stored, b.DiskGUID)
}

// ---------------------------------------------------------------------------
// Plumbing.
// ---------------------------------------------------------------------------

func emitJSON(v any) int {
	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	if err := enc.Encode(v); err != nil {
		fmt.Fprintf(os.Stderr, "partitionguard: encoding JSON: %v\n", err)
		return exitUsage
	}
	return 0
}

func usage(w io.Writer) {
	fmt.Fprint(w, `partitionguard - read-only MBR / GPT partition table parser

USAGE
  partitionguard <command> <image> [options]

COMMANDS
  info      Detect the partitioning scheme and report GPT header fields,
            including header and partition-array CRC32 validation.
  list      List every partition: start/end LBA, sector count, size and type.
  verify    Run structural checks (CRC32s, primary vs backup GPT header,
            partition overlap, partitions past end of image, MBR signature).
  help      Show this message.

OPTIONS
  --json               Emit machine-readable JSON instead of text.
  --sector-size <n>    Logical sector size in bytes (default 512).
  -h, --help           Show this message.

EXIT STATUS
  0   success (for verify: every check passed or was skipped)
  1   usage error, or the image could not be read or parsed
  2   verify only: at least one check FAILED

NOTES
  partitionguard never writes to the image it inspects; it opens the file
  read-only and performs bounds-checked reads.

EXAMPLES
  partitionguard info disk.img
  partitionguard list disk.img --json
  partitionguard verify disk.img
  partitionguard list disk.img --sector-size 4096
`)
}

func isHelpArg(s string) bool {
	switch s {
	case "-h", "--help", "help", "-help", "--h":
		return true
	}
	return false
}

func main() {
	os.Exit(run(os.Args[1:]))
}

func run(argv []string) int {
	if len(argv) == 0 {
		// Double-clicked in Explorer rather than run from a prompt: ask the
		// one question the program needs and stay on screen. Printing usage
		// and exiting here is what made the window vanish instantly.
		if interactiveConsole() {
			runGuided()
			return 0
		}
		usage(os.Stderr)
		return exitUsage
	}
	if isHelpArg(argv[0]) {
		usage(os.Stdout)
		return 0
	}

	cmd := argv[0]
	switch cmd {
	case "info", "list", "verify":
	default:
		fmt.Fprintf(os.Stderr, "partitionguard: unknown command %q\n\n", cmd)
		usage(os.Stderr)
		return exitUsage
	}

	args := reorderFlags(argv[1:], map[string]bool{"sector-size": true})

	fs := flag.NewFlagSet(cmd, flag.ContinueOnError)
	fs.SetOutput(io.Discard)
	asJSON := fs.Bool("json", false, "emit JSON")
	sectorSize := fs.Int("sector-size", 512, "logical sector size")
	if err := fs.Parse(args); err != nil {
		if err == flag.ErrHelp {
			usage(os.Stdout)
			return 0
		}
		fmt.Fprintf(os.Stderr, "partitionguard: %v\n\n", err)
		usage(os.Stderr)
		return exitUsage
	}

	rest := fs.Args()
	if len(rest) == 0 {
		fmt.Fprintf(os.Stderr, "partitionguard: %s requires an image path\n\n", cmd)
		usage(os.Stderr)
		return exitUsage
	}
	if len(rest) > 1 {
		fmt.Fprintf(os.Stderr, "partitionguard: %s takes exactly one image path (got %d: %s)\n\n",
			cmd, len(rest), strings.Join(rest, ", "))
		usage(os.Stderr)
		return exitUsage
	}
	if !validSectorSize(*sectorSize) {
		fmt.Fprintf(os.Stderr, "partitionguard: invalid --sector-size %d (must be a power of two between 512 and 65536)\n", *sectorSize)
		return exitUsage
	}

	path := rest[0]
	im, err := openImage(path, *sectorSize)
	if err != nil {
		fmt.Fprintf(os.Stderr, "partitionguard: %v\n", err)
		return exitUsage
	}
	defer im.Close()

	if abs, aerr := filepath.Abs(path); aerr == nil {
		im.path = abs
	}

	l, err := analyze(im)
	if err != nil {
		fmt.Fprintf(os.Stderr, "partitionguard: %s: %v\n", path, err)
		return exitUsage
	}
	if l.scheme == "none" && cmd != "verify" {
		fmt.Fprintf(os.Stderr, "partitionguard: %s: no partition table found (no 0x55AA MBR signature at offset 510 and no %q signature at LBA 1)\n",
			path, gptSignature)
		return exitUsage
	}

	switch cmd {
	case "info":
		return cmdInfo(l, *asJSON)
	case "list":
		return cmdList(l, *asJSON)
	default:
		return cmdVerify(l, *asJSON)
	}
}

func validSectorSize(n int) bool {
	if n < 512 || n > 65536 {
		return false
	}
	return n&(n-1) == 0
}
