hk = require('luahooks32core')

-- Memory protection
-- Allow writing to code areas
local PAGE_EXECUTE_READWRITE = 0x40
function hk.protectRWX(addr, len)
	return hk.protect(addr, len, PAGE_EXECUTE_READWRITE)
end

hk._openImages = hk._openImages or {}
-- open library memoization
function hk.open(name)
	if name==nil then name = '_baseImage' end
	local img = hk._openImages[name]
	if img then return img end
	img = hk._openRaw(name~='_baseImage' and name or nil)
	hk._openImages[name] = img
	return img
end

-- Scanning
-- Convert text-style scan pattern into code-style pattern
-- Input: '11 22 33 ? 44'
-- Output: '\x11\x22\x33\x00\x44', 'xxx?x'
local function patToCode(p)
	if p:find('^str:') then -- raw string
		local pat = p:sub(4, #p)
		local mask = ('x'):rep(#pat)
		return pat, mask
	else -- hex pattern
		if p:find('[^a-fA-F0-9 \r\n\t%?]') then
			error('hk pattern: pattern contains invalid character', 3) end
		local patT, maskT = {}, {}
		for word in p:gmatch('[^ \r\n\t]+') do
			if word:find('%?') then
				table.insert(patT, string.char(0))
				table.insert(maskT, '?')
			else
				local val = tonumber(word, 16)
				assert(val and val>=0 and val<=255, 'invalid word in scan pattern: '..word, 3)
				table.insert(patT, string.char(val))
				table.insert(maskT, 'x')
			end
		end
		return table.concat(patT), table.concat(maskT)
	end
end
-- Scan
-- hk.scan('15 1 ? 1f') - return first match (or nil)
-- hk.scan('15 1 ? 1f', 1) - return nth match (or nil) (starting from 1)
-- hk.scan('15 1 ? 1f', true) - return list of all matches (or {})
local unhookAll, rehookAll
hk._scanResultCache = hk._scanResultCache or {}
function hk.scan(pat, opt, img)
	if type(pat)~='string' then
		error('hk.scan: argument #1: expected string', 2) end
	local code, mask = patToCode(pat)
	img = img or hk.open()
	opt = opt or 1
	
	local _
	local cacheEntry = tostring(img)..':'..pat..':'..tostring(opt)
	local res = hk._scanResultCache[cacheEntry]
	if res then return res end
	
	if opt==true then -- find all matches
		res = {}
		unhookAll()
		while true do
			local suc,a = pcall(hk._scanRaw, img, code, mask, #res)
			if not a then break end
			table.insert(res, a)
		end
		rehookAll()
	elseif type(opt)=='number' and opt%1==0 and opt>0 then -- find nth match
		unhookAll()
		suc,res = pcall(hk._scanRaw, img, code, mask, opt-1)
		rehookAll()
	else
		error('hk.scan: argument #2: expected true, positive integer, or nil', 2)
	end
	
	hk._scanResultCache[cacheEntry] = res
	return res
end


-- Writing
local function hexToStr(h)
	local t = {}
	if h:find('[^a-zA-Z0-9 \r\n\t]') then
		error('hk.write: invalid character in hex string', 3) end
	for w in h:gmatch('[^ \r\n\t]+') do
		local v = tonumber(w, 16)
		if not (v and v>=0 and v<=255) then
			error('hk.write: invalid hex number: '..w, 3) end
		table.insert(t, string.char(v))
	end
	return table.concat(t)
end
local customWriters = {
	hex = function(addr, str, len)
		local data = hexToStr(str)
		if len then return addr+#data end
		return hk.writeStr(addr, data)
	end,
	str = function(addr, data, len)
		if len then return addr+#data end
		return hk.writeStr(addr, data)
	end,
	char = function(addr, val, len)
		if len then return addr+1 end
		return hk.writeChar(addr, val)
	end,
	short = function(addr, val, len)
		if len then return addr+2 end
		return hk.writeShort(addr, val)
	end,
	int = function(addr, val, len)
		if len then return addr+4 end
		return hk.writeInt(addr, val)
	end,
	rel = function(addr, val, len)
		if len then return addr+4 end
		return hk.writeInt(addr, val - (addr+4))
	end,
	float = function(addr, val, len)
		if len then return addr+4 end
		return hk.writeFloat(addr, val)
	end,
	double = function(addr, val, len)
		if len then return addr+8 end
		return hk.writeFloat(addr, val)
	end,
}
local customWriterDefaults = {
	number = 'int',
	string = 'hex',
}
local function writeData(addr, data, typ, len)
	if type(data)=='table' then
		if typ then
			error('hk.write: argument #3: expected nil when argument #2 is table', 2) end
		if not table.islist(data) then
			error('hk.write: argument #2: table must be a list', 2) end
		
		local ntyp = nil
		for i,v in ipairs(data) do
			if type(v)=='string' and v:sub(#v,#v)==':' then
				ntyp = v:sub(1,#v-1)
				if not customWriters[ntyp] then
					error('hk.write: argument #2: expected writer type at index '
						..i..', got \''..ntyp..'\'', 2) end
			else
				addr = hk.write(addr, data[i], ntyp, len)
				ntyp = nil
			end
		end
		return addr
	else
		if not typ then
			typ = customWriterDefaults[type(data)]
			if not typ then
				error('hk.write: argument #2: expected string, number, or table') end
		end
		if not customWriters[typ] then
			error('hk.write: argument #3: expected writer type, got \''..typ..'\'') end
		return customWriters[typ](addr, data, len)
	end
end
function hk.write(addr, data, typ)
	return writeData(addr, data, typ, false)
end
-- write to a write-protected area by turning off write protection first,
-- then re-enable write protection afterward
function hk.patch(addr, data, typ)
	local len = writeData(addr, data, typ, true) - addr
	local oldProt = hk.protectRWX(addr, len)
	local addrW = writeData(addr, data, typ, false)
	hk.protect(addr, len, oldProt)
	return addrW
end


-- Hooking
local function writeTrampoline(trAddr, hkAddr, regsPtr)
	return hk.write(trAddr, {
		-- save registers and flags
		'a3',regsPtr,                           -- mov [regsPtr],eax
		'b8',regsPtr,                           -- mov eax,regsPtr
		'89 58 04',                             -- mov [eax+0x04],ebx
		'89 48 08',                             -- mov [eax+0x08],ecx
		'89 50 0c',                             -- mov [eax+0x0c],edx
		'89 70 10',                             -- mov [eax+0x10],esi
		'89 78 14',                             -- mov [eax+0x14],edi
		'89 60 18',                             -- mov [eax+0x18],esp
		'89 68 1c',                             -- mov [eax+0x1c],ebp
		'c7 40 20',hkAddr,                      -- mov [eax+0x20],hkAddr (res->eip)
		'9c',                                   -- pushfd
		'5a',                                   -- pop edx
		'89 50 24',                             -- mov [eax+0x24],edx (regs->flags)
		'c7 40 28',0,                           -- mov [eax+0x28],0 (regs->brk = 0)
		-- load arguments into ecx+edx and call
		'89 c1',                                -- mov ecx,eax
		'ba',hk._getLuaStatePtr(),        -- mov edx,L
		'e8','rel:',hk._getCallbackPtr(), -- call hook function
		-- restore registers
		'b8',regsPtr,                           -- mov eax,regsPtr
		'8b 58 04',                             -- mov ebx,[eax+0x04]
		'8b 48 08',                             -- mov ecx,[eax+0x08]
		'8b 50 0c',                             -- mov edx,[eax+0x0c]
		'8b 70 10',                             -- mov esi,[eax+0x10]
		'8b 78 14',                             -- mov edi,[eax+0x14]
		'8b 60 18',                             -- mov esp,[eax+0x18]
		'8b 68 1c',                             -- mov ebp,[eax+0x1c]
		-- if regs.brk, restore eax and retn; otherwise continue
		'83 78 28 00',                          -- cmp dword ptr [eax+0x28],0
		'74 06',                                -- je (past eax restore and retn)
		'a1',regsPtr,                           -- mov eax,[regsPtr]
		'c3',                                   -- retn
		-- restore flags and eax
		'8b 50 24',                             -- mov edx,[eax+0x24] (regs->flags)
		'52',                                   -- push edx
		'9d',                                   -- popfd
		'a1',regsPtr,                           -- mov eax,[regsPtr]
		-- after this, moved code will be written, followed by a jump back
	})
end
local function writeTrampolineReturn(trAddr, retAddr)
	return hk.write(trAddr, {
		'e9','rel:',retAddr,                    -- jmp retAddr
	})
end
local regsStructSize = 4*11
local regsStruct = {
	eax= 0, ebx= 4, ecx= 8, edx=12,
	esi=16, edi=20, esp=24, ebp=28,
	eip=32, flags=36, brk=40,
}
local regsList = {'eax','ebx','ecx','edx','esi','edi','esp','ebp','eip','flags','brk'}
local function newRegs()
	return hk.malloc(regsStructSize)
end
local function readRegsStruct(regsPtr)
	local regs = {}
	for _,regname in ipairs(regsList) do
		regs[regname] = hk.readInt(regsPtr + regsStruct[regname])
	end
	return regs
end

-- Basic instruction length determination for automatic trampoline construction
-- Add to this table as needed to define instructions
-- a value of true indicates a position-dependent instruction that cannot be moved
local instrLen = {
	['1b'] = {
		['c9'] = 2, -- sbb ecx,ecx
	},
	['23'] = {
		['c8'] = 2, -- and ecx,eax
	},
	['33'] = {
		['c4'] = 2, -- xor eax,esp
	},
	['3b'] = {
		['15'] = 6, -- cmp edx,i32
	},
	['50'] = 1, -- push eax
	['51'] = 1, -- push ecx
	['53'] = 1, -- push ebx
	['55'] = 1, -- push ebp
	['56'] = 1, -- push esi
	['57'] = 1, -- push edi
	['5d'] = 1, -- pop ebp
	['5e'] = 1, -- pop esi
	['5f'] = 1, -- pop edi
	['68'] = 5, -- push i32
	['6a'] = 2, -- push i8
	['72'] = true, -- jb rel8
	['74'] = true, -- jz rel8
	['75'] = true, -- jnz rel8
	['81'] = {
		['ec'] = 6, -- sub esp,i32
		['c2'] = 6, -- add edx,i32
	},
	['83'] = {
		['c4'] = 3, -- add esp,i8
		['e4'] = 3, -- and esp,i8
		['ec'] = 3, -- sub esp,i8
	},
	['85'] = {
		['c0'] = 2, -- test eax,eax
	},
	['89'] = {
		['0d'] = 6, -- mov [i32],ecx
		['15'] = 6, -- mov [i32],edx
	},
	['8b'] = {
		['0d'] = 6, -- mov ecx,i32
		['40'] = 3, -- mov eax,[eax+i8]
		['44'] = {
			['24'] = 4, -- mov eax,[esp+i8]
		},
		['45'] = 3, -- mov eax,[ebp+i8]
		['6b'] = 3, -- mov ebp,[ebx+i8]
		['c8'] = 2, -- mov ecx,eax
		['dc'] = 2, -- mov ebx,esp
		['e5'] = 2, -- mov esp,ebp
		['ec'] = 2, -- mov ebp,esp
		['f1'] = 2, -- mov esi,ecx
	},
	['8d'] = {
		['34'] = {
			['01'] = 3, -- lea esi,[ecx+eax]
		},
		['90'] = 6, -- lea edx,[eax+i32]
	},
	['a1'] = 5, -- mov eax,i32
	['b8'] = 5, -- mov eax,i32
	['c3'] = 1, -- retn
	['d9'] = {
		['47'] = 3, -- fld:32 [edi+i8]
		['87'] = 6, -- fld:32 [edi+i32]
	},
	['dd'] = {
		['5c'] = {
			['24'] = 4, -- fstp:64 [esp+i8]
		},
	},
	['e8'] = true, -- call rel32
	['f7'] = {
		['d9'] = 2, -- neg ecx
	},
}
local function readByteHex(addr)
	return ('%02x'):format(hk.readChar(addr)%256)
end
local jmpoutLen = 5 -- Length of the long jump instruction to be inserted as a hook
-- Determine the minimum code length >= jmpoutLen that can be
-- copied out into the trampoline.
-- Returns false if unrecognized instructions
-- Returns true if instructions are position-dependent and cannot be moved
local function determineOverwriteLen(addr)
	local len = 0
	while len < jmpoutLen do
		local val = readByteHex(addr)
		local il = instrLen[val]
		local ofs = 0
		while type(il)=='table' do
			ofs = ofs+1
			val = readByteHex(addr+ofs)
			il = il[val]
		end
		if not il then return false end
		if il==true then return true end
		addr = addr + il
		len = len + il
	end
	return len
end
-- Write into existing function to jump to trampoline
local function writeJumpout(rh)
	local oldProt = hk.protectRWX(rh.hkAddr, rh.hkLen)
	local hkAddrW = hk.write(rh.hkAddr, {
		'e9','rel:',rh.trAddr, -- jmp trAddr
	})
	for i = 1, rh.hkLen-jmpoutLen do -- fill extra space with nops
		hkAddrW = hk.write(hkAddrW, '90')
	end
	hk.protect(rh.hkAddr, rh.hkLen, oldProt) -- restore write protection on jumpout
	return hkAddrW
end
-- Write code copied from hooked function into trampoline
local function writeOldCode(rh)
	local oldProt = hk.protectRWX(rh.hkAddr, rh.hkLen)
	hk.writeStr(rh.hkAddr, rh.movedCode)
	hk.protect(rh.hkAddr, rh.hkLen, oldProt)
end

hk._registeredHooks = hk._registeredHooks or {} -- map of addr -> list of callbacks
-- Remove/replace all hooks, used when scanning
unhookAll = function()
	for _,rh in pairs(hk._registeredHooks) do
		writeOldCode(rh)
	end
end
rehookAll = function()
	for _,rh in pairs(hk._registeredHooks) do
		writeJumpout(rh)
	end
end
-- Called from C in trampoline
function _bllua_hk_callback(regsPtr)
	local regs = readRegsStruct(regsPtr)
	local rh = hk._registeredHooks[regs.eip]
	if not rh then
		error('_bllua_hk_callback: no callback registered at address') end
	rh.callback(regs)
end
-- Main raw hooking function
function hk.hook(hkAddr, callback, hkLen)
	if type(hkAddr)~='number' or hkAddr<0 or hkAddr%1~=0 then
		error('hk.hook: argument #1: expected number >0 integer', 2) end
	if type(callback)~='function' then
		error('hk.hook: argument #2: expected function', 2) end
	if hkLen~=nil and (type(hkLen)~='number' or hkLen<0 or hkLen%1~=0) then
		error('hk.hook: argument #3: expected nil or number', 2) end
	
	-- if a hook is already registered, overwrite it
	-- todo: multiple hooks? package names?
	if hk._registeredHooks[hkAddr] then
		print('hk.hook: warning: a hook is already registered at address '..
			('%08x'):format(hkAddr)..', will overwrite.')
		hk._registeredHooks[hkAddr].callback = callback
		return
	end
	
	if hkLen then
		if hkLen<jmpoutLen then
			error('hk.hook: argument #3: overwrite length must be >= '
				..jmpoutLen, 2) end
	else
		hkLen = determineOverwriteLen(hkAddr)
		if hkLen==false then
			error('hk.hook: could not automatically determine instruction length. '
				..'please specify a length >= '..jmpoutLen..' in argument #3', 2)
		elseif hkLen==true then
			error('hk.hook: the hook location contains position-dependent code! '
				..'please move the hook or use a different hooking method', 2)
		end
	end
	
	-- create register save struct
	local regsPtr = newRegs()
	
	-- create trampoline code
	local movedCode = hk.readStr(hkAddr, hkLen)
	local trLen = 256 + #movedCode -- eh, good enough
	local trAddr = hk.malloc(trLen)
	local trOldProt = hk.protectRWX(trAddr, trLen) -- allow execution
	local trAddrW = trAddr
	trAddrW = writeTrampoline(trAddrW, hkAddr, regsPtr)
	trAddrW = hk.writeStr(trAddrW, movedCode)
	trAddrW = writeTrampolineReturn(trAddrW, hkAddr + hkLen)
	
	-- create info struct
	local rh = {
		hkAddr = hkAddr,
		hkLen = hkLen,
		regsPtr = regsPtr,
		trAddr = trAddr,
		trLen = trLen,
		trOldProt = trOldProt,
		movedCode = movedCode,
		callback = callback,
	}
	
	-- save to hook registry
	hk._registeredHooks[hkAddr] = rh
	
	-- write jump out
	writeJumpout(rh)
end
-- Remove a hook made by hk.hookRaw
function hk.unhook(hkAddr)
	if type(hkAddr)~='number' or hkAddr<0 or hkAddr%1~=0 then
		error('hk.unhook: argument #1: expected number >0 integer', 2) end
	local rh = hk._registeredHooks[hkAddr]
	assert(rh, 'hk.unhook: no hook registered at address '..
		('%08x'):format(hkAddr))
	
	-- remove jumpout
	writeOldCode(rh)
	
	-- delete allocated data
	hk.protect(rh.trAddr, rh.trLen, rh.trOldProt)
	hk.free(rh.trAddr)
	hk.free(rh.regsPtr)
	
	-- remove from hook registry
	hk._registeredHooks[hkAddr] = nil
end
function hex(v) return v and ('%08x'):format(v):sub(1,8) end
-- called when blocklua unloads
function hk.unhookAll()
	unhookAll()
	for _,rh in pairs(hk._registeredHooks) do
		hk.protect(rh.trAddr, rh.trLen, rh.trOldProt)
		hk.free(rh.trAddr)
		hk.free(rh.regsPtr)
	end
	hk._registeredHooks = {}
end
-- Utility to display addresses as hex
function hk.hex(v)
	return ('%08x'):format(v)
end

-- todo: stack manipulation


_bllua_on_unload['libhooks'] = hk.unhookAll

return hk


-- testing
--[[

'dofile('modules/lualib/luahooks32.lua')

'f = hk.scan('55 8B EC 83 E4 C0 83 EC 38 56 57 8B 7D 08 8B 47 44 D1 E8 A8 01 74 5C 8B 47 28', 2)

'hk.hook(f, function(regs) print(regs) end)


not documented:
hk.open
hk.write / hk.patch formats other than hex
hk.protect
hk.protectRWX

]]