package kernelcache

//go:generate go tool stringer -type=SubsystemStart -output mig_string.go

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"strings"

	"github.com/apex/log"
	"github.com/blacktop/arm64-cgo/disassemble"
	"github.com/blacktop/go-macho"
	"github.com/blacktop/go-macho/types"
	"github.com/blacktop/ipsw/pkg/disass"
)

type SubsystemStart uint32

const (
	arcade_register_subsystem  SubsystemStart = 0xC90F
	catch_exc_subsystem        SubsystemStart = 0x961
	catch_mach_exc_subsystem   SubsystemStart = 0x965
	clock_subsystem            SubsystemStart = 0x3E8
	host_priv_subsystem        SubsystemStart = 0x190
	is_iokit_subsystem         SubsystemStart = 0xAF0
	mach_eventlink_subsystem   SubsystemStart = 0xAEDA8
	mach_host_subsystem        SubsystemStart = 0xC8
	mach_port_subsystem        SubsystemStart = 0xC80
	mach_vm_subsystem          SubsystemStart = 0x12c0
	mach_voucher_subsystem     SubsystemStart = 0x1518
	memory_entry_subsystem     SubsystemStart = 0x1324
	processor_set_subsystem    SubsystemStart = 0xFA0
	processor_subsystem        SubsystemStart = 0xBB8
	task_restartable_subsystem SubsystemStart = 0x1F40
	task_subsystem             SubsystemStart = 0xD48
	thread_act_subsystem       SubsystemStart = 0xE10
	UNDReply_subsystem         SubsystemStart = 0x1838
	vm32_map_subsystem         SubsystemStart = 0xED8
)

type MigHash struct {
	Num           int
	KObjIdx       int
	KRoutine      uint64 /* Kernel server routine */
	KReplySize    uint32 /* Size of kernel reply msg */
	KReplyDescCnt uint32 /* Number of descs in kernel reply msg */
}

type MachMsgHeader struct {
	MsghBits        uint32
	MsghSize        uint32
	MsghRemotePort  uint32
	MsghLocalPort   uint32
	MsghVoucherPort uint32
	MsghId          uint32
}

type KernRoutineDescriptor struct {
	ImplRoutine     uint64 /* Server work func pointer */
	KStubRoutine    uint64 /* Unmarshalling func pointer */
	ArgC            uint32 /* Number of argument words */
	DescrCount      uint32 /* Number complex descriptors */
	ReplyDescrCount uint32 /* Number descriptors in reply */
	MaxReplyMsg     uint32 /* Max size for reply msg */
}

type migKernSubsystemHdr struct {
	KServer  uint64         /* pointer to kernel demux routine */
	Start    SubsystemStart /* Min routine number */
	End      uint32         /* Max routine number + 1 */
	Maxsize  uint32         /* Max reply message size */
	_        uint32         // padding
	Reserved uint64         /* reserved for MIG use */

}

type MigKernSubsystem struct {
	migKernSubsystemHdr
	Routines []KernRoutineDescriptor /* Kernel routine descriptor array */
}

func (m MigKernSubsystem) LookupRoutineName(idx int) string {
	switch m.Start {
	case mach_vm_subsystem:
		if idx >= len(machVmSubsystemFuncs) {
			return "<unknown>"
		} else {
			return machVmSubsystemFuncs[idx]
		}
	case mach_port_subsystem:
		if idx >= len(machPortSubsystemFuncs) {
			return "<unknown>"
		} else {
			return machPortSubsystemFuncs[idx]
		}
	case mach_host_subsystem:
		if idx >= len(machHostSubsystemFuncs) {
			return "<unknown>"
		} else {
			return machHostSubsystemFuncs[idx]
		}
	case host_priv_subsystem:
		if idx >= len(hostPrivSubsystemFuncs) {
			return "<unknown>"
		} else {
			return hostPrivSubsystemFuncs[idx]
		}
	case clock_subsystem:
		if idx >= len(clockSubsystemFuncs) {
			return "<unknown>"
		} else {
			return clockSubsystemFuncs[idx]
		}
	case processor_subsystem:
		if idx >= len(processorSubsystemFuncs) {
			return "<unknown>"
		} else {
			return processorSubsystemFuncs[idx]
		}
	case processor_set_subsystem:
		if idx >= len(processorSetSubsystemFuncs) {
			return "<unknown>"
		} else {
			return processorSetSubsystemFuncs[idx]
		}
	case is_iokit_subsystem:
		if idx >= len(isIokitSubsystemProcessorSetSubsystemFuncs) {
			return "<unknown>"
		} else {
			return isIokitSubsystemProcessorSetSubsystemFuncs[idx]
		}
	case task_subsystem:
		if idx >= len(taskSubsystemFuncs) {
			return "<unknown>"
		} else {
			return taskSubsystemFuncs[idx]
		}
	case thread_act_subsystem:
		if idx >= len(threadActSubsystemFuncs) {
			return "<unknown>"
		} else {
			return threadActSubsystemFuncs[idx]
		}
	case vm32_map_subsystem:
		if idx >= len(vm32MapSubsystemFuncs) {
			return "<unknown>"
		} else {
			return vm32MapSubsystemFuncs[idx]
		}
	case UNDReply_subsystem:
		if idx >= len(undReplySubsystemFuncs) {
			return "<unknown>"
		} else {
			return undReplySubsystemFuncs[idx]
		}
	case mach_voucher_subsystem:
		if idx >= len(machVoucherSubsystemFuncs) {
			return "<unknown>"
		} else {
			return machVoucherSubsystemFuncs[idx]
		}
	case memory_entry_subsystem:
		if idx >= len(memoryEntrySubsystemFuncs) {
			return "<unknown>"
		} else {
			return memoryEntrySubsystemFuncs[idx]
		}
	case task_restartable_subsystem:
		if idx >= len(taskRestartableSubsystemFuncs) {
			return "<unknown>"
		} else {
			return taskRestartableSubsystemFuncs[idx]
		}
	case catch_exc_subsystem:
		if idx >= len(catchExcSubsystemFuncs) {
			return "<unknown>"
		} else {
			return catchExcSubsystemFuncs[idx]
		}
	case catch_mach_exc_subsystem:
		if idx >= len(catch_mach_exc_subsystem_Funcs) {
			return "<unknown>"
		} else {
			return catch_mach_exc_subsystem_Funcs[idx]
		}
	case arcade_register_subsystem:
		if idx >= len(arcade_register_subsystem_Funcs) {
			return "<unknown>"
		} else {
			return arcade_register_subsystem_Funcs[idx]
		}
	case mach_eventlink_subsystem:
		if idx >= len(mach_eventlink_subsystem_Funcs) {
			return "<unknown>"
		} else {
			return mach_eventlink_subsystem_Funcs[idx]
		}
	default:
		return "<unknown>"
	}
}

func (m MigKernSubsystem) String() string {
	var out strings.Builder
	out.WriteString(fmt.Sprintf("%s: %s\t%s=%d %s=%d %s=%d\n",
		colorAddr("%#x", m.KServer),
		colorSubSystem(m.Start.String()),
		colorField("start"), m.Start,
		colorField("end"), m.End,
		colorField("max_sz"), m.Maxsize,
	))
	for idx, r := range m.Routines {
		if r.KStubRoutine == 0 {
			continue // skip empty routines
		}
		out.WriteString(fmt.Sprintf("    %s: ", colorAddr("%#x", r.KStubRoutine)))
		out.WriteString(colorBold(m.LookupRoutineName(idx)))
		out.WriteString(fmt.Sprintf("\t%s=%02d %s=%#x %s=%02d %s=%d %s=%d %s=%d\n",
			colorName("num"), idx+int(m.Start),
			colorName("impl"), r.ImplRoutine,
			colorName("argc"), r.ArgC,
			colorName("descr"), r.DescrCount,
			colorName("reply_descr"), r.ReplyDescrCount,
			colorName("max_reply_msg"), r.MaxReplyMsg,
		))
	}
	return out.String()
}

func getMigInitFunc(m *macho.File) (*types.Function, error) {
	var ref uint64

	cstrs, err := m.GetCStrings()
	if err != nil {
		return nil, err
	}

	found := false
	for str, addr := range cstrs["__TEXT.__cstring"] {
		if strings.Contains(str, "mig_e") {
			ref = addr
			found = true
			break
		}
	}
	if !found {
		return nil, fmt.Errorf("failed to find 'mig_e' anchor")
	}

	text := m.Section("__TEXT_EXEC", "__text")
	if text == nil {
		return nil, fmt.Errorf("failed to find __TEXT_EXEC.__text section")
	}
	data, err := text.Data()
	if err != nil {
		return nil, fmt.Errorf("failed to get data from __TEXT_EXEC.__text section: %v", err)
	}

	engine := disass.NewMachoDisass(m, &disass.Config{
		Data:         data,
		StartAddress: text.Addr,
		Quiet:        true,
	})
	if err := engine.Triage(); err != nil {
		return nil, fmt.Errorf("first pass triage failed: %v", err)
	}

	if ok, loc := engine.Contains(ref); ok {
		migInit, err := m.GetFunctionForVMAddr(loc)
		if err != nil {
			return nil, fmt.Errorf("failed to get function with xref to 'mig_e' anchor: %v", err)
		}
		return &migInit, nil
	}

	return nil, fmt.Errorf("failed to find 'mig_init' address")
}

func migOperandRegister(instr *disassemble.Inst, idx int) (disassemble.Register, bool) {
	if instr == nil || int(instr.NumOps) <= idx || instr.Operands[idx].NumRegisters == 0 {
		return disassemble.REG_NONE, false
	}
	return instr.Operands[idx].Registers[0], true
}

func migOperandImmediate(instr *disassemble.Inst, idx int) (uint64, bool) {
	if instr == nil || int(instr.NumOps) <= idx {
		return 0, false
	}
	return instr.Operands[idx].Immediate, true
}

func getMigE(r *bytes.Reader, migInit *types.Function) (uint64, error) {
	var migE uint64

	var instrValue uint32
	var decoder disassemble.Decoder
	var prevInstr disassemble.Inst
	var hasPrev bool

	startAddr := migInit.StartAddr

	for {
		err := binary.Read(r, binary.LittleEndian, &instrValue)
		if err != nil {
			if err == io.EOF {
				break
			}
			return 0, fmt.Errorf("failed to read instruction @ %#x: %v", startAddr, err)
		}

		var instruction disassemble.Inst
		if err := decoder.DecomposeInto(startAddr, instrValue, &instruction); err != nil {
			startAddr += uint64(binary.Size(uint32(0)))
			continue
		}

		if disass.IsLoadLiteral(&instruction) {
			if imm, ok := migOperandImmediate(&instruction, 1); ok {
				migE = imm
				break
			}
		} else if hasPrev && prevInstr.Operation == disassemble.ARM64_ADRP &&
			(instruction.Operation == disassemble.ARM64_ADD ||
				instruction.Operation == disassemble.ARM64_LDR ||
				instruction.Operation == disassemble.ARM64_LDRB ||
				instruction.Operation == disassemble.ARM64_LDRSW) {
			adrpRegister, ok := migOperandRegister(&prevInstr, 0)
			if ok {
				adrpImm, ok := migOperandImmediate(&prevInstr, 1)
				if ok {
					srcRegister, ok := migOperandRegister(&instruction, 1)
					if ok {
						validPattern := true
						if adrpRegister == srcRegister {
							switch instruction.Operation {
							case disassemble.ARM64_LDR, disassemble.ARM64_LDRB, disassemble.ARM64_LDRSW:
								if imm, ok := migOperandImmediate(&instruction, 1); ok {
									adrpImm += imm
								} else {
									validPattern = false
								}
							case disassemble.ARM64_ADD:
								if imm, ok := migOperandImmediate(&instruction, 2); ok {
									adrpImm += imm
								} else {
									validPattern = false
								}
							}
						}
						if validPattern {
							migE = adrpImm
							break
						}
					}
				}
			}
		}
		// fmt.Printf("%#08x:  %s\t%s\n", uint64(startAddr), disassemble.GetOpCodeByteString(instrValue), instruction)

		prevInstr = instruction
		hasPrev = true
		startAddr += uint64(binary.Size(uint32(0)))
	}

	if migE == 0 {
		return 0, fmt.Errorf("failed to find 'mig_e' table address in mig_init")
	}

	return migE, nil
}

// getMigSubsystemPointers discovers the list of mig_subsystem pointers in the
// mig_e[] table by walking pointer-sized entries starting at migEAddr and
// validating that each entry looks like a reasonable mig_subsystem header
// within __DATA_CONST.__const. This avoids relying on brittle instruction
// patterns (such as CMP-immediate) in mig_init, which can change between
// kernel versions.
func getMigSubsystemPointers(m *macho.File, migEAddr uint64) ([]uint64, error) {
	const maxEntries = 64 // generous upper bound for number of MIG subsystems

	dataConst := m.Section("__DATA_CONST", "__const")
	if dataConst == nil {
		return nil, macho.ErrMachOSectionNotFound
	}
	dataConstData, err := dataConst.Data()
	if err != nil {
		return nil, err
	}

	dataConstStart := dataConst.Addr
	dataConstEnd := dataConstStart + uint64(len(dataConstData))

	hdrSize := uint64(binary.Size(migKernSubsystemHdr{}))

	var subsystems []uint64

	for i := range uint64(maxEntries) {
		ptr, err := m.GetPointerAtAddress(migEAddr + i*uint64(binary.Size(uint64(0))))
		if err != nil {
			// If we fail on the very first entry, something is wrong with migEAddr.
			if i == 0 {
				return nil, fmt.Errorf("failed to read mig_e entry at %#x: %v", migEAddr, err)
			}
			break
		}

		// A NULL entry is treated as the end of the table if we've already
		// collected some subsystems.
		if ptr == 0 {
			if len(subsystems) == 0 {
				continue
			}
			break
		}

		// mig_subsystem headers live in __DATA_CONST.__const.
		if ptr < dataConstStart || ptr+hdrSize > dataConstEnd {
			// If we haven't collected any valid entries yet, keep scanning in
			// case the first few entries are something else.
			if len(subsystems) == 0 {
				continue
			}
			break
		}

		off := ptr - dataConstStart
		r := bytes.NewReader(dataConstData[off:])

		var hdr migKernSubsystemHdr
		if err := binary.Read(r, binary.LittleEndian, &hdr); err != nil {
			if len(subsystems) == 0 {
				return nil, fmt.Errorf("failed to read mig subsystem header at %#x: %v", ptr, err)
			}
			break
		}

		// Basic sanity checks on the header:
		if hdr.Start == 0 || hdr.End <= uint32(hdr.Start) {
			if len(subsystems) == 0 {
				continue
			}
			break
		}
		if hdr.End-uint32(hdr.Start) > 0x1000 {
			if len(subsystems) == 0 {
				continue
			}
			break
		}

		subsystems = append(subsystems, ptr)
	}

	if len(subsystems) == 0 {
		return nil, fmt.Errorf("failed to infer mig subsystem table size from mig_e (addr=%#x)", migEAddr)
	}

	return subsystems, nil
}

func GetMigSubsystems(m *macho.File) (uint64, []MigKernSubsystem, error) {
	if m.FileTOC.FileHeader.Type == types.MH_FILESET {
		var err error
		m, err = m.GetFileSetFileByName("com.apple.kernel")
		if err != nil {
			return 0, nil, fmt.Errorf("failed to get fileset entry 'com.apple.kernel': %v", err)
		}
	}

	migInit, err := getMigInitFunc(m)
	if err != nil {
		return 0, nil, err
	}
	data, err := m.GetFunctionData(*migInit)
	if err != nil {
		return 0, nil, err
	}

	migEAddr, err := getMigE(bytes.NewReader(data), migInit)
	if err != nil {
		return 0, nil, err
	}

	subsystems, err := getMigSubsystemPointers(m, migEAddr)
	if err != nil {
		return migEAddr, nil, err
	}

	log.WithField("mig_kern_subsystem table", fmt.Sprintf("%#x", migEAddr)).Infof("Found")

	dataConst := m.Section("__DATA_CONST", "__const")
	if dataConst == nil {
		return migEAddr, nil, macho.ErrMachOSectionNotFound
	}
	dataConstData, err := dataConst.Data()
	if err != nil {
		return migEAddr, nil, err
	}

	r := bytes.NewReader(dataConstData)

	var migs []MigKernSubsystem

	for i := range subsystems {
		r.Seek(int64(subsystems[i]-dataConst.Addr), io.SeekStart)

		var mig MigKernSubsystem
		if err := binary.Read(r, binary.LittleEndian, &mig.migKernSubsystemHdr); err != nil {
			return migEAddr, nil, err
		}
		mig.migKernSubsystemHdr.KServer = m.SlidePointer(mig.migKernSubsystemHdr.KServer)
		mig.Routines = make([]KernRoutineDescriptor, mig.End-uint32(mig.Start))
		if err := binary.Read(r, binary.LittleEndian, &mig.Routines); err != nil {
			return migEAddr, nil, err
		}
		for i, routine := range mig.Routines {
			routine.ImplRoutine = m.SlidePointer(routine.ImplRoutine)
			routine.KStubRoutine = m.SlidePointer(routine.KStubRoutine)
			mig.Routines[i] = routine
		}

		migs = append(migs, mig)
	}

	return migEAddr, migs, nil
}