#    This file is part of Metasm, the Ruby assembly manipulation suite
#    Copyright (C) 2006-2009 Yoann GUILLOT
#
#    Licence is LGPL, see LICENCE in the top-level directory


#
# This sample defines an ApiHook class, that you can subclass to easily hook functions
# in a debugged process. Your custom function will get called whenever an API function is,
# giving you access to the arguments, you can also take control just before control returns
# to the caller.
# See the example in the end for more details.
# As a standalone application, it hooks WriteFile in the 'notepad' process, and make it
# skip the first two bytes of the buffer.
#

require 'metasm'

class ApiHook
	# rewrite this function to list the hooks you want
	# return an array of hashes
	def setup
		#[{ :function => 'WriteFile', :abi => :stdcall },	# standard function hook
		# { :module => 'Foo.dll', :rva => 0x2433,		# arbitrary code hook
		#   :abi => :fastcall, :hookname => 'myhook' }]		# hooks named pre_myhook/post_myhook
	end

	# initialized from a Debugger or a process description that will be debugged
	# sets the hooks up, then run_forever
	def initialize(dbg)
		if not dbg.kind_of? Metasm::Debugger
			process = Metasm::OS.current.find_process(dbg)
			raise 'no such process' if not process
			dbg = process.debugger
		end
		dbg.loadallsyms
		@dbg = dbg
		setup.each { |h| setup_hook(h) }
		init_prerun if respond_to?(:init_prerun)	# allow subclass to do stuff before main loop
		@dbg.run_forever
	end

	# setup one function hook
	def setup_hook(h)
		pre  =  "pre_#{h[:hookname] || h[:function]}"
		post = "post_#{h[:hookname] || h[:function]}"

		@nargs = h[:nargs] || method(pre).arity if respond_to?(pre)

		if target = h[:address]
		elsif target = h[:rva]
			modbase = @dbg.modulemap[h[:module]]
			raise "cant find module #{h[:module]} in #{@dbg.modulemap.join(', ')}" if not modbase
			target += modbase[0]
		else
			target = h[:function]
		end

		@dbg.bpx(target) {
			catch(:finish) {
				@cur_abi = h[:abi]
				@ret_longlong = h[:ret_longlong]
				if respond_to? pre
					args = read_arglist
					send pre, *args
				end
				if respond_to? post
					@dbg.bpx(@dbg.func_retaddr, true) {
						retval = read_ret
						send post, retval, args
					}
				end
			}
		}
	end

	# retrieve the arglist at func entry, from @nargs & @cur_abi
	def read_arglist
		nr = @nargs
		args = []

		if (@cur_abi == :fastcall or @cur_abi == :thiscall) and nr > 0
			args << @dbg.get_reg_value(:ecx)
			nr -= 1
		end

		if @cur_abi == :fastcall and nr > 0
			args << @dbg.get_reg_value(:edx)
			nr -= 1
		end

		nr.times { |i| args << @dbg.func_arg(i) }

		args
       	end

	# retrieve the function returned value
	def read_ret
		ret = @dbg.func_retval
		if @ret_longlong
			ret = (ret & 0xffffffff) | (@dbg[:edx] << 32)
		end
		ret
	end

	# patch the value of an argument
	# only valid in pre_hook
	# nr starts at 0
	def patch_arg(nr, value)
		case @cur_abi
		when :fastcall
			case nr
			when 0
				@dbg.set_reg_value(:ecx, value)
				return
			when 1
				@dbg.set_reg_value(:edx, value)
				return
			else
				nr -= 2
			end
		when :thiscall
			case nr
			when 0
				@dbg.set_reg_value(:ecx, value)
				return
			else
				nr -= 1
			end
		end

		@dbg.func_arg_set(nr, value)
	end

	# patch the function return value
	# only valid post_hook
	def patch_ret(val)
		if @ret_longlong
			@dbg.set_reg_value(:edx, (val >> 32) & 0xffffffff)
			val &= 0xffffffff
		end
		@dbg.func_retval_set(val)
	end

	# skip the function call
	# only valid in pre_hook
	def finish(retval)
		patch_ret(retval)
		@dbg.ip = @dbg.func_retaddr
		case @cur_abi
		when :fastcall
			@dbg[:esp] += 4*(@nargs-2) if @nargs > 2
		when :thiscall
			@dbg[:esp] += 4*(@nargs-1) if @nargs > 1
		when :stdcall
			@dbg[:esp] += 4*@nargs
		end
		@dbg.sp += @dbg.cpu.size/8
		throw :finish
	end
end



if __FILE__ == $0

# this is the class you have to define to hook
# 
# setup() defines the list of hooks as an array of hashes
# for exported functions, simply use :function => function name
# for arbitrary hook, :module => 'module.dll', :rva => 0x1234, :hookname => 'myhook' (call pre_myhook/post_myhook)
# :abi can be :stdcall (windows standard export), :fastcall or :thiscall, leave empty for cdecl
# if pre_<function> is defined, it is called whenever the function is entered, via a bpx (int3)
# if post_<function> is defined, it is called whenever the function exists, with a bpx on the return address setup at func entry
# the pre_hook receives all arguments to the original function
#  change them with patch_arg(argnr, newval)
#  read memory with @dbg.memory_read_int(ptr), or @dbg.memory[ptr, length]
#  skip the function call with finish(fake_retval) (!) needs a correct :abi & param count !
# the post_hook receives the function return value
#  change it with patch_ret(newval)
class MyHook < ApiHook
	def setup
		[{ :function => 'WriteFile', :abi => :stdcall }]
	end

	def init_prerun
		puts "hooks ready, save a file in notepad"
	end

	def pre_WriteFile(handle, pbuf, size, pwritten, overlap)
		# we can skip the function call with this
		#finish(28)

		# spy on the api / trace calls
		bufdata = @dbg.memory[pbuf, size]
		puts "writing #{bufdata.inspect}"

		# but we can also mess with the args
		# ex: skip first 2 bytes of the buffer
		patch_arg(1, pbuf+2)
		patch_arg(2, size-2)
	end

	def post_WriteFile(retval, arglistcopy)
		# we can patch the API return value with this
		#patch_retval(42)

		# finish messing with the args: fake the nrofbyteswritten
		handle, pbuf, size, pwritten, overlap = arglistcopy
		written = @dbg.memory_read_int(pwritten)
		if written == size
			# if written everything, patch the value so that the program dont detect our intervention
			@dbg.memory_write_int(pwritten, written+2)
		end

		puts "write retval: #{retval}, written: #{written} bytes"
	end
end

# name says it all
Metasm::WinOS.get_debug_privilege

# run our Hook engine on a running 'notepad' instance
MyHook.new('notepad')

# the script ends when notepad exits

end