package winpwn

import (
	"unsafe"
)

// DllCharacteristics bits (winnt.h IMAGE_DLLCHARACTERISTICS_*).
const (
	dllCharHighEntropyVA  = 0x0020
	dllCharDynamicBase    = 0x0040 // ASLR
	dllCharForceIntegrity = 0x0080
	dllCharNXCompat       = 0x0100 // DEP
	dllCharNoIsolation    = 0x0200
	dllCharNoSEH          = 0x0400
	dllCharAppContainer   = 0x1000
	dllCharGuardCF        = 0x4000 // CFG
)

// Data directory indices (winnt.h IMAGE_DIRECTORY_ENTRY_*).
const (
	dirEntrySecurity      = 4 // Authenticode; VirtualAddress here is a *file offset*, not an RVA.
	dirEntryLoadConfig    = 10
	dirEntryComDescriptor = 14 // .NET CLR header
)

// IMAGE_GUARD_CF_INSTRUMENTED, from the GuardFlags field of the Load Config
// Directory: set when the binary actually has CFG checks emitted, as
// opposed to just the (necessary but not sufficient) DllCharacteristics bit.
const imageGuardCFInstrumented = 0x00000100

// CheckSecResult mirrors pwntools'/checksec's binary protection summary,
// adapted to the mitigations that actually exist on PE/Windows.
type CheckSecResult struct {
	Is64Bit bool

	ASLR           bool // IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE
	HighEntropyVA  bool // IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA (64-bit ASLR range)
	DEP            bool // IMAGE_DLLCHARACTERISTICS_NX_COMPAT
	CFG            bool // Control Flow Guard
	ForceIntegrity bool
	IsolationAware bool
	AppContainer   bool
	DotNET         bool // has a CLR/COM descriptor header, i.e. is a managed binary

	// SEH/SafeSEH only mean anything for 32-bit PE32 images: x64 uses
	// table-based structured exception handling and isn't subject to the
	// classic SEH-chain-overwrite technique at all.
	SEHApplicable bool
	HasSEH        bool // false if compiled with the /SAFESEH-equivalent IMAGE_DLLCHARACTERISTICS_NO_SEH
	SafeSEH       bool // SEHandlerTable present in Load Config

	GSHeuristic bool // SecurityCookie present in Load Config (best-effort, see Checksec doc comment)

	AuthenticodeSigned bool // IMAGE_DIRECTORY_ENTRY_SECURITY present (presence only, not cryptographically verified)
}

// Checksec inspects compile-time/link-time exploit mitigations, the Go
// analogue of pwntools'/winchecksec's binary checksec report.
//
// GSHeuristic is exactly that: a heuristic. Unlike ASLR/DEP/CFG which are
// global, unambiguous flags, /GS stack-cookie insertion is decided by the
// compiler per function. The presence of a non-zero SecurityCookie slot in
// the Load Config Directory only tells you the binary *could* use stack
// cookies, not that the specific function you're exploiting does — verify
// against the actual disassembly before relying on it.
func (p *PEFile) Checksec() (*CheckSecResult, error) {
	h, err := p.header()
	if err != nil {
		return nil, err
	}

	r := &CheckSecResult{
		Is64Bit:        h.is64,
		ASLR:           h.dllCharacteristics&dllCharDynamicBase != 0,
		HighEntropyVA:  h.dllCharacteristics&dllCharHighEntropyVA != 0,
		DEP:            h.dllCharacteristics&dllCharNXCompat != 0,
		CFG:            h.dllCharacteristics&dllCharGuardCF != 0,
		ForceIntegrity: h.dllCharacteristics&dllCharForceIntegrity != 0,
		IsolationAware: h.dllCharacteristics&dllCharNoIsolation == 0,
		AppContainer:   h.dllCharacteristics&dllCharAppContainer != 0,
		SEHApplicable:  !h.is64,
		HasSEH:         h.dllCharacteristics&dllCharNoSEH == 0,
	}

	r.DotNET = h.dataDirectory[dirEntryComDescriptor].VirtualAddress != 0
	r.AuthenticodeSigned = h.dataDirectory[dirEntrySecurity].VirtualAddress != 0

	lc, err := p.readLoadConfig(h)
	if err == nil && lc != nil {
		r.SafeSEH = !h.is64 && lc.sehHandlerTable != 0
		r.GSHeuristic = lc.securityCookie != 0
		// Corroborate the DllCharacteristics CFG bit with the GuardFlags
		// instrumentation bit when we have a Load Config to check it against.
		r.CFG = r.CFG && lc.guardFlags&imageGuardCFInstrumented != 0
	}

	return r, nil
}

// loadConfig64/loadConfig32 mirror winnt.h's IMAGE_LOAD_CONFIG_DIRECTORY64/32:
// same field order in both, only pointer-sized members change width. Named
// (not anonymous) so unsafe.Offsetof can validate against each directory's
// self-reported Size — older toolchains emit a shorter struct with no
// SafeSEH/Guard CF fields at all, and reading past Size would misattribute
// zeroed padding as "feature present".
type loadConfig64 struct {
	Size                           uint32
	TimeDateStamp                  uint32
	MajorVersion                   uint16
	MinorVersion                   uint16
	GlobalFlagsClear               uint32
	GlobalFlagsSet                 uint32
	CriticalSectionDefaultTimeout  uint32
	DeCommitFreeBlockThreshold     uint64
	DeCommitTotalFreeThreshold     uint64
	LockPrefixTable                uint64
	MaximumAllocationSize          uint64
	VirtualMemoryThreshold         uint64
	ProcessAffinityMask            uint64
	ProcessHeapFlags               uint32
	CSDVersion                     uint16
	DependentLoadFlags             uint16
	EditList                       uint64
	SecurityCookie                 uint64
	SEHandlerTable                 uint64
	SEHandlerCount                 uint64
	GuardCFCheckFunctionPointer    uint64
	GuardCFDispatchFunctionPointer uint64
	GuardCFFunctionTable           uint64
	GuardCFFunctionCount           uint64
	GuardFlags                     uint32
}

type loadConfig32 struct {
	Size                           uint32
	TimeDateStamp                  uint32
	MajorVersion                   uint16
	MinorVersion                   uint16
	GlobalFlagsClear               uint32
	GlobalFlagsSet                 uint32
	CriticalSectionDefaultTimeout  uint32
	DeCommitFreeBlockThreshold     uint32
	DeCommitTotalFreeThreshold     uint32
	LockPrefixTable                uint32
	MaximumAllocationSize          uint32
	VirtualMemoryThreshold         uint32
	ProcessAffinityMask            uint32
	ProcessHeapFlags               uint32
	CSDVersion                     uint16
	DependentLoadFlags             uint16
	EditList                       uint32
	SecurityCookie                 uint32
	SEHandlerTable                 uint32
	SEHandlerCount                 uint32
	GuardCFCheckFunctionPointer    uint32
	GuardCFDispatchFunctionPointer uint32
	GuardCFFunctionTable           uint32
	GuardCFFunctionCount           uint32
	GuardFlags                     uint32
}

// Field offsets, computed by the compiler instead of hand-counted, used to
// validate against each Load Config's self-reported Size.
var (
	offsetOf64SecurityCookie = uint32(unsafe.Offsetof(loadConfig64{}.SecurityCookie))
	offsetOf64SEHandlerCount = uint32(unsafe.Offsetof(loadConfig64{}.SEHandlerCount))
	offsetOf64GuardFlags     = uint32(unsafe.Offsetof(loadConfig64{}.GuardFlags))

	offsetOf32SecurityCookie = uint32(unsafe.Offsetof(loadConfig32{}.SecurityCookie))
	offsetOf32SEHandlerCount = uint32(unsafe.Offsetof(loadConfig32{}.SEHandlerCount))
	offsetOf32GuardFlags     = uint32(unsafe.Offsetof(loadConfig32{}.GuardFlags))
)

// loadConfigInfo holds the handful of Load Config Directory fields checksec
// cares about, already normalized to a common width.
type loadConfigInfo struct {
	securityCookie  uint64
	sehHandlerTable uint64
	guardFlags      uint32
}

// readLoadConfig parses the Load Config Directory, respecting its own Size
// field so we never trust fields beyond what the linker actually emitted.
func (p *PEFile) readLoadConfig(h peHeader) (*loadConfigInfo, error) {
	dir := h.dataDirectory[dirEntryLoadConfig]
	if dir.VirtualAddress == 0 {
		return nil, nil
	}
	offset := p.RVAToFileOffset(dir.VirtualAddress)
	if offset == 0 {
		return nil, nil
	}

	info := &loadConfigInfo{}
	if h.is64 {
		var lc loadConfig64
		if err := p.readStructAt(offset, &lc); err != nil {
			return nil, err
		}
		if lc.Size > offsetOf64SecurityCookie {
			info.securityCookie = lc.SecurityCookie
		}
		if lc.Size > offsetOf64SEHandlerCount {
			info.sehHandlerTable = lc.SEHandlerTable
		}
		if lc.Size > offsetOf64GuardFlags {
			info.guardFlags = lc.GuardFlags
		}
		return info, nil
	}

	var lc loadConfig32
	if err := p.readStructAt(offset, &lc); err != nil {
		return nil, err
	}
	if lc.Size > offsetOf32SecurityCookie {
		info.securityCookie = uint64(lc.SecurityCookie)
	}
	if lc.Size > offsetOf32SEHandlerCount {
		info.sehHandlerTable = uint64(lc.SEHandlerTable)
	}
	if lc.Size > offsetOf32GuardFlags {
		info.guardFlags = lc.GuardFlags
	}
	return info, nil
}