-- Implement the marker table.
--
local markerObjects = {}
_G.Markers = { Clear = function() markerObjects = {} end }
setmetatable(_G.Markers, {
	__index = function(_, name)
		local marker = markerObjects[name]
		return marker and Rule.ReadMarker(marker)
	end
})

-- Parses a table of rules and returns a table of valid Rule objects.
--
local function parseRuleList(arg)
   local result = {}
	local it = 1
	while true do
		local v = arg[it]
      if v == nil then break end
		local t = type(v)

		if t == "string" then
			if v:sub(1,1) == "#" then
				v = v:sub(2)
				local m = markerObjects[v]
				if not m then
					m = Rule.Marker()
					markerObjects[v] = m
				end
				result[#result+1] = m
			elseif type(arg[it+1]) == "number" then
				result[#result+1] = Rule.Instruction(v, arg[it+1])
				it = it + 1
			else
				result[#result+1] = Rule.Pattern(v)
			end
		elseif t == "table" then
			result[#result+1] = Rule.Pad(v[1], v[2])
		elseif t == "function" then
			result[#result+1] = Rule.Filter(v)
		elseif t == "userdata" then
			result[#result+1] = v
		else
			error("Did not expect: " .. t)
		end
		it = it + 1
	end
	return result
end

-- Implement Rule.All and Rule.Any.
--
function Rule.All(arg)
	return Rule.AllImpl(parseRuleList(arg))
end
function Rule.Any(arg)
	return Rule.AnyImpl(parseRuleList(arg))
end

-- Hex helper, override print.
--
function Hex(value) return value and string.format("%x", value) or "nil" end
function Addr(value) return value and string.format("%x", value + Image.base) or "nil" end

function print(...)
	local args = table.pack(...)
	local msg = ""
	for i=1, args.n do
		if #msg ~= 0 then
			msg = msg .. " " .. tostring(args[i])
		else
			msg = tostring(args[i])
		end
	end
	Log(msg)
end

local function findAny(patternList)
	local output = {0}
	local n = 0
	for k,v in ipairs(patternList) do
		for rva in Image:Find(Rule.Pattern(v[1])) do
			local res = rva + ((#v[1]+1)/3) + v[2]
			output[n+2] = res
			n = n + 1
		end
	end
	output[1] = n
	return output
end

-- Implement searchers for user call gadgets.
--
Register("_fast_return_gadgets", R_ARR, function()
	return findAny{
		{ "E8 ?? 0F 05 00", -3 }, -- call rel32
		{ "0F 8? 0F 05 00", -3 }, -- jcc rel32
		{ "FF 15 ?? 0F 05", -2 }, -- call [rip+rel32]
		{ "48 8D ?? ?? 0F 05", -2 }, -- lea [rip+rel32]
		{ "48 89 ?? ?? 0F 05", -2 }, -- mov [rip+rel32]
		{ "4C 8D ?? ?? 0F 05", -2 }, -- lea [rip+rel32]
		{ "4C 89 ?? ?? 0F 05", -2 }, -- mov [rip+rel32]
	} 
end)

Register("_fix_arg_gadgets", R_ARR, function()
	return findAny{
		{ "F3 41 0F 59 C3", -2 }, -- mulss, mulps, vmulss, vmulps
		{ "F3 43 0F 59 C3", -2 },
		{ "41 0F 59 C3", -2 },
		{ "F3 0F 59 C3", -2 },
		{ "C4 C1 ?? 59 C3", -2 }
	}
end)

Register("_stack_spoof_gadgets", R_ARR, function()
	-- Matches a call m64 instruction
	-- reg  = /2
	local callPattern = { "FF 10", "FF 38" }
	
	-- Matches an epilogue
	--
	local epiPatterns = {
		{ "48 83 C4 08 C3",          "FF FF FF 8F FF" },         -- [add rsp, imm8],  no negative/misaligned immediates
		{ "48 81 C4 08 00 00 00 C3", "FF FF FF 0F FF FF FF FF" } -- [add rsp, imm32], no negative/misaligned/huge immediates
	}

	-- For each epilogue pattern:
	--
	local results = {}
	for _,v in ipairs(epiPatterns) do
		-- For each call instruction suffix length:
		--
		for suf=5,0,-1 do
			-- Find the gadget:
			--
			local value = callPattern[1] .. string.rep(" 00", suf) .. " " .. v[1]
			local mask =  callPattern[2] .. string.rep(" 00", suf) .. " " .. v[2]
			for match in Image:Find(Rule.Pattern(value, mask)) do
				-- If valid:
				--
				if Image:DecodeLen(match) == (2+suf) then
					-- Decode the sub instruction, if it is not at arg area, write the result:
					--
					local subinsn = Image:Decode(match+2+suf)
					if subinsn.imm0 >= (8*9) then
						results[#results+1] = {
							callRva = match,
							subRva =  match + 2 + suf,
							delta =   subinsn.imm0 + 8
						}
					end
				end
			end
		end
	end

	-- Sort the table based on call complexity.
	--
	table.sort(results, function(a, b) return (a.subRva-a.callRva)>(b.subRva-b.callRva) end)

	-- Get table length, trimming to a maximum of 32 and write the immediates.
	--
	local len = math.min(32, #results)
	local output = {}
	output[1] = len
	for i=1,len do
		output[i*2 + 0] = results[i].subRva
		output[i*2 + 1] = results[i].delta
	end
	return output
end)