package extract

import (
	"encoding/binary"
	"sort"

	"golang.org/x/arch/x86/x86asm"

	"ringer/pefile"
)

// confidenceRank orders confidence labels so dedupe keeps the strongest signal.
var confidenceRank = map[string]int{"high": 3, "medium": 2, "low": 1}

// ExtractIoctls returns every IOCTL code found in the driver, deduplicated and
// sorted by code. It uses two passes:
//
//  1. Disassembly of .text — immediates in CMP/MOV are the canonical IOCTL
//     dispatch pattern (high/medium confidence).
//  2. Data scan of .rdata/.data/.pdata — 32-bit values that decode to a
//     plausible IOCTL and are referenced by code (medium confidence).
func ExtractIoctls(f *pefile.File) []IoctlCode {
	best := map[uint32]IoctlCode{}
	xrefs := map[uint32]bool{}

	if f.Machine.IsX86() {
		if text := f.Section(".text"); text != nil && len(text.Data) > 0 {
			mode := 32
			if f.Is64 {
				mode = 64
			}
			disasmPass(f, text.Data, text.VirtualAddress, mode, best, xrefs)
		}
	}

	for _, name := range []string{".rdata", ".data", ".pdata"} {
		if s := f.Section(name); s != nil && len(s.Data) > 0 {
			dataPass(s.Data, s.VirtualAddress, best, xrefs)
		}
	}

	// Dispatch tables (arrays of IOCTLs) are iterated in a loop, so their
	// entries are not xref'd by CMP and are missed by the xref-gated data pass.
	for _, name := range []string{".rdata", ".data"} {
		if s := f.Section(name); s != nil && len(s.Data) > 0 {
			tablePass(s.Data, s.VirtualAddress, best)
		}
	}

	out := make([]IoctlCode, 0, len(best))
	for _, c := range best {
		out = append(out, c)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Code < out[j].Code })
	return out
}

// addCandidate records a candidate IOCTL, keeping the highest-confidence entry
// for each code.
func addCandidate(best map[uint32]IoctlCode, c IoctlCode) {
	existing, ok := best[c.Code]
	if !ok || confidenceRank[c.Confidence] > confidenceRank[existing.Confidence] {
		best[c.Code] = c
	}
}

// decodedInst is a single decoded instruction with its RVA.
type decodedInst struct {
	inst x86asm.Inst
	rva  uint32
}

// disasmPass linearly sweeps a code section, extracting 32-bit immediates from
// comparison/load instructions, collecting data cross-references, and detecting
// switch jump-table dispatch.
func disasmPass(f *pefile.File, data []byte, baseRVA uint32, mode int, best map[uint32]IoctlCode, xrefs map[uint32]bool) {
	var insts []decodedInst
	for off := 0; off < len(data); {
		inst, err := x86asm.Decode(data[off:], mode)
		if err != nil {
			// Desync (data embedded in .text, jump table, padding): skip one
			// byte and retry so a single bad byte does not abort the sweep.
			off++
			continue
		}
		insts = append(insts, decodedInst{inst, baseRVA + uint32(off)})
		off += inst.Len
	}

	for _, di := range insts {
		switch di.inst.Op {
		case x86asm.CMP:
			collectXrefs(di.inst, di.rva, mode, xrefs)
			extractImmediates(di.inst, di.rva, "high", best)
		case x86asm.MOV:
			extractImmediates(di.inst, di.rva, "medium", best)
		}
	}

	switchPass(f, insts, best)
}

// collectXrefs records data addresses referenced by CMP memory operands (the
// IOCTL table-dispatch pattern cmp IoControlCode, [table+idx*4]).
func collectXrefs(inst x86asm.Inst, instRVA uint32, mode int, xrefs map[uint32]bool) {
	for _, arg := range inst.Args {
		mem, ok := arg.(x86asm.Mem)
		if !ok {
			continue
		}
		if mode == 64 && mem.Base == x86asm.RIP {
			target := int64(instRVA) + int64(inst.Len) + int64(int32(mem.Disp))
			if target >= 0 {
				xrefs[uint32(target)] = true
			}
		} else if mode == 32 && mem.Base == 0 && mem.Index == 0 && mem.Disp > 0 {
			xrefs[uint32(mem.Disp)] = true
		}
	}
}

// switchPass detects compiler-generated switch jump-table dispatch and extracts
// every case value. MSVC lowers a switch on the IOCTL code to:
//
//	sub/add reg, imm32   ; index = code - BASE (ADD uses 2^32 - BASE)
//	cmp reg, imm         ; range check
//	ja/jae default       ; unsigned bounds check
//	lea reg, [rip+bt]    ; byte-index table (sparse switches)
//	movzx reg, [reg+reg] ; case = bt[index]
//	lea reg, [rip+jt]    ; jump table
//	movsxd reg, [reg+4*reg]
//
// The base IOCTL is imm32 (SUB) or 2^32 - imm32 (ADD). For a sparse switch the
// byte-index table maps each index in [0, range] to a case number; the default
// case fills the holes, so every index whose case differs from the default is a
// real IOCTL (base + index).
func switchPass(f *pefile.File, insts []decodedInst, best map[uint32]IoctlCode) {
	for i := 0; i+2 < len(insts); i++ {
		reg, imm, ok := subAddImm(insts[i].inst)
		if !ok {
			continue
		}
		creg, ok := cmpRegImm(insts[i+1].inst)
		if !ok || creg != reg {
			continue
		}
		j := insts[i+2].inst
		if j.Op != x86asm.JA && j.Op != x86asm.JAE {
			continue
		}
		var base uint32
		if insts[i].inst.Op == x86asm.SUB {
			base = uint32(imm)
		} else {
			base = uint32(-imm)
		}
		rangeImm, ok := insts[i+1].inst.Args[1].(x86asm.Imm)
		if !ok {
			continue
		}
		maxIndex := int64(rangeImm)
		if maxIndex < 0 || maxIndex > 4096 {
			recordSwitchIoctl(best, base, insts[i].rva)
			continue
		}

		// A sparse switch has a byte-index table; a dense switch indexes the
		// jump table directly. Only the sparse form lets us enumerate cases, so
		// fall back to the base value when the table is absent or unreadable.
		byteTable, ok := findByteIndexTable(insts, i+3)
		if !ok {
			recordSwitchIoctl(best, base, insts[i].rva)
			continue
		}
		bt := f.ReadRVA(byteTable, int(maxIndex)+1)
		if bt == nil {
			recordSwitchIoctl(best, base, insts[i].rva)
			continue
		}
		def := mostFrequentByte(bt)
		for idx, c := range bt {
			if c == def {
				continue
			}
			recordSwitchIoctl(best, base+uint32(idx), insts[i].rva)
		}
	}
}

// recordSwitchIoctl records a switch-derived IOCTL candidate at high confidence.
func recordSwitchIoctl(best map[uint32]IoctlCode, code uint32, rva uint32) {
	ok, note := plausibleIoctl(code)
	if !ok {
		return
	}
	dt, fn, m, a := DecodeIoctl(code)
	addCandidate(best, IoctlCode{
		Code:       code,
		DeviceType: dt,
		Function:   fn,
		Method:     m,
		Access:     a,
		Confidence: "high",
		Source:     "switch",
		RVA:        rva,
		Note:       note,
	})
}

// findByteIndexTable locates the byte-index table of a sparse switch. Two MSVC
// lowerings exist:
//
//	PowerStrip:   LEA reg, [rip+bt] ; MOVZX ... [reg+reg]        (LEA then MOVZX)
//	ThrottleStop: LEA reg, [rip+base] ; ... saves ... ; MOVZX ... [reg+idx+disp]
//
// The second separates the LEA from the MOVZX with register saves, so locate the
// MOVZX first and resolve its base register backward to the LEA. The byte-index
// table is then base + disp.
func findByteIndexTable(insts []decodedInst, start int) (uint32, bool) {
	for k := start; k < len(insts) && k < start+16; k++ {
		if insts[k].inst.Op != x86asm.MOVZX {
			continue
		}
		base, disp, ok := movzxMem(insts[k].inst)
		if !ok {
			continue
		}
		for j := k - 1; j >= start && j >= k-8; j-- {
			if insts[j].inst.Op != x86asm.LEA {
				continue
			}
			reg, target, ok := leaReg(insts[j].inst, insts[j].rva)
			if !ok || reg != base {
				continue
			}
			return uint32(int64(target) + disp), true
		}
	}
	return 0, false
}

// movzxMem returns the base register and displacement of a MOVZX memory operand
// that indexes a table (base + index registers both present). This is the
// byte-index-table load of a sparse switch.
func movzxMem(inst x86asm.Inst) (x86asm.Reg, int64, bool) {
	if inst.Op != x86asm.MOVZX {
		return 0, 0, false
	}
	mem, ok := inst.Args[1].(x86asm.Mem)
	if !ok {
		return 0, 0, false
	}
	if mem.Base == 0 || mem.Index == 0 {
		return 0, 0, false
	}
	return mem.Base, int64(int32(mem.Disp)), true
}

// leaReg returns the destination register and RIP-relative target of a
// LEA reg, [RIP+disp] instruction.
func leaReg(inst x86asm.Inst, instRVA uint32) (x86asm.Reg, uint32, bool) {
	if inst.Op != x86asm.LEA {
		return 0, 0, false
	}
	reg, ok := inst.Args[0].(x86asm.Reg)
	if !ok {
		return 0, 0, false
	}
	target, ok := leaRIP(inst, instRVA)
	if !ok {
		return 0, 0, false
	}
	return reg, target, true
}

// leaRIP returns the RIP-relative target of a LEA reg, [RIP+disp] instruction.
func leaRIP(inst x86asm.Inst, instRVA uint32) (uint32, bool) {
	if inst.Op != x86asm.LEA {
		return 0, false
	}
	mem, ok := inst.Args[1].(x86asm.Mem)
	if !ok {
		return 0, false
	}
	if mem.Base != x86asm.RIP {
		return 0, false
	}
	// mem.Disp is stored raw (uint32) for disp32, so sign-extend to recover
	// negative RIP-relative displacements (e.g. 0xffffe0ce -> -7986).
	target := int64(instRVA) + int64(inst.Len) + int64(int32(mem.Disp))
	if target < 0 {
		return 0, false
	}
	return uint32(target), true
}

// mostFrequentByte returns the byte value that occurs most often in b. In a
// sparse-switch byte-index table this is the default case, which fills the holes
// between the real cases.
func mostFrequentByte(b []byte) byte {
	var freq [256]int
	for _, c := range b {
		freq[c]++
	}
	var best byte
	bestCount := -1
	for c := 0; c < 256; c++ {
		if freq[c] > bestCount {
			bestCount = freq[c]
			best = byte(c)
		}
	}
	return best
}

// subAddImm returns the destination register and 32-bit immediate of a
// SUB/ADD reg, imm32 instruction. inst.Args is a fixed [6]Arg array with nil
// padding, so unused operands are checked by type assertion, not length.
func subAddImm(inst x86asm.Inst) (x86asm.Reg, int64, bool) {
	if inst.Op != x86asm.SUB && inst.Op != x86asm.ADD {
		return 0, 0, false
	}
	reg, ok := inst.Args[0].(x86asm.Reg)
	if !ok {
		return 0, 0, false
	}
	imm, ok := inst.Args[1].(x86asm.Imm)
	if !ok {
		return 0, 0, false
	}
	return reg, int64(imm), true
}

// cmpRegImm reports whether inst is CMP reg, imm and returns the register.
func cmpRegImm(inst x86asm.Inst) (x86asm.Reg, bool) {
	if inst.Op != x86asm.CMP {
		return 0, false
	}
	reg, ok := inst.Args[0].(x86asm.Reg)
	if !ok {
		return 0, false
	}
	if _, ok := inst.Args[1].(x86asm.Imm); !ok {
		return 0, false
	}
	return reg, true
}

// extractImmediates pulls 32-bit immediate operands from an instruction and
// records any that decode to a plausible IOCTL.
func extractImmediates(inst x86asm.Inst, instRVA uint32, conf string, best map[uint32]IoctlCode) {
	for _, arg := range inst.Args {
		imm, ok := arg.(x86asm.Imm)
		if !ok {
			continue
		}
		iv := int64(imm)
		// Only 32-bit immediates are IOCTL candidates. Sign-extended negatives
		// (high bit set) and small positives both fall in this range; 64-bit
		// immediates fall outside it.
		if iv < -0x80000000 || iv > 0xFFFFFFFF {
			continue
		}
		code := uint32(iv)
		ok, note := plausibleIoctl(code)
		if !ok {
			continue
		}
		dt, fn, m, a := DecodeIoctl(code)
		addCandidate(best, IoctlCode{
			Code:       code,
			DeviceType: dt,
			Function:   fn,
			Method:     m,
			Access:     a,
			Confidence: conf,
			Source:     "disasm",
			RVA:        instRVA,
			Note:       note,
		})
	}
}

// dataPass scans a data section for 32-bit values that decode to a plausible
// IOCTL AND whose address is referenced by code. Unreferenced DWORDs are too
// noisy (any 32-bit value decodes to *some* IOCTL), so they are skipped.
func dataPass(data []byte, baseRVA uint32, best map[uint32]IoctlCode, xrefs map[uint32]bool) {
	for off := 0; off+4 <= len(data); off += 4 {
		rva := baseRVA + uint32(off)
		if !xrefs[rva] {
			continue
		}
		code := binary.LittleEndian.Uint32(data[off : off+4])
		recordDataIoctl(best, code, rva, "medium", "data")
	}
}

// tablePass scans a data section for runs of consecutive DWORDs that all decode
// to plausible IOCTLs sharing the same custom vendor device type (0x8000-0xBFFF).
// Such a run is a plain IOCTL array, which is iterated in a loop and therefore
// not xref'd by CMP (so the xref-gated data pass misses it).
//
// The vendor-device-type requirement is what keeps this from firing on the two
// common lookalikes: tables of RVAs (device type 0x0001-0x00FF) and UTF-16
// strings (device type 0x0020-0x007E). Third-party drivers must use device types
// >= 0x8000, so a run of three DWORDs sharing such a type is a real table, not
// coincidence.
//
// {ioctl, handler} pair tables (8/16-byte stride) are not scanned: the strided
// scan fires on too many non-IOCTL tables in the 0x8000-0xBFFF range to be
// reliable, so those are left to the disasm pass.
func tablePass(data []byte, baseRVA uint32, best map[uint32]IoctlCode) {
	for off := 0; off+4 <= len(data); off += 4 {
		code := binary.LittleEndian.Uint32(data[off : off+4])
		ok, _ := plausibleIoctl(code)
		if !ok {
			continue
		}
		dt, _, _, _ := DecodeIoctl(code)
		if dt < 0x8000 {
			continue
		}
		run := []uint32{code}
		for o := off + 4; o+4 <= len(data); o += 4 {
			c := binary.LittleEndian.Uint32(data[o : o+4])
			ok, _ := plausibleIoctl(c)
			if !ok {
				break
			}
			cdt, _, _, _ := DecodeIoctl(c)
			if cdt != dt {
				break
			}
			run = append(run, c)
		}
		if len(run) < 3 {
			continue
		}
		for i, c := range run {
			recordDataIoctl(best, c, baseRVA+uint32(off+i*4), "low", "table")
		}
		off += (len(run) - 1) * 4
	}
}

// recordDataIoctl decodes a raw 32-bit value and records it as an IOCTL
// candidate if it is plausible.
func recordDataIoctl(best map[uint32]IoctlCode, code uint32, rva uint32, conf, source string) {
	ok, note := plausibleIoctl(code)
	if !ok {
		return
	}
	dt, fn, m, a := DecodeIoctl(code)
	addCandidate(best, IoctlCode{
		Code:       code,
		DeviceType: dt,
		Function:   fn,
		Method:     m,
		Access:     a,
		Confidence: conf,
		Source:     source,
		RVA:        rva,
		Note:       note,
	})
}