require 'msf/core'

module Msf

###
#
# This module exposes methods for talking to WDBRPC daemons
#
###
module Exploit::Remote::WDBRPC_Client

	include Exploit::Remote::WDBRPC
	include Auxiliary::Report

	attr_accessor :wdbrpc_info, :udp_sock

	def initialize(info = {})
		super
		register_options(
			[
				Opt::RHOST,
				Opt::RPORT(17185),
			], Msf::Exploit::Remote::WDBRPC_Client)
	end


	def wdbrpc_client_connect
		self.wdbrpc_info = {}

		wdbrpc_client_disconnect()

		self.udp_sock = Rex::Socket::Udp.create(
			{
				'Context' => {'Msf' => framework, 'MsfExploit' => self}
			}
		)
		add_socket(self.udp_sock)

		wdbrpc_client_send_disconnect()

		udp_sock.sendto(wdbrpc_request_connect(rhost), rhost, rport, 0)
		res,src = udp_sock.recvfrom(65535, 5)
		if not res
			print_error("No response to TARGET_CONNECT (WDB4)")
			return
		end


		if res.length > 0 and res.length < 80
			print_status("#{rhost}: Unknown response: '#{res.unpack("H*")[0]}'")
			return
		end

		if res.empty?
			print_error("#{rhost}: No response from the target")
			return
		end

		self.wdbrpc_info = wdbrpc_parse_connect_reply(res)
		print_status("#{rhost} Connected to #{self.wdbrpc_info[:rt_vers]} - #{self.wdbrpc_info[:rt_bsp_name]} (#{self.wdbrpc_info[:rt_bootline]})")

		report_note(
			:host   => rhost,
			:port   => rport,
			:proto  => 'udp',
			:type   => 'vxworks.target_info',
			:data   => res,
			:update => :unique
		)
	end

	def wdbrpc_client_connect2
		self.wdbrpc_info = {}

		wdbrpc_client_disconnect()

		self.udp_sock = Rex::Socket::Udp.create(
			{
				'Context' => {'Msf' => framework, 'MsfExploit' => self}
			}
		)
		add_socket(self.udp_sock)

		wdbrpc_client_send_disconnect()

		udp_sock.sendto(wdbrpc_request_connect2(rhost), rhost, rport, 0)
		res,src = udp_sock.recvfrom(65535, 5)
		if not res
			print_error("No response to TARGET_CONNECT2")
			return
		end

		if res.length < 80
			print_status("#{rhost}: Unknown response: '#{res.unpack("H*")[0]}'")
			return
		end

		self.wdbrpc_info = wdbrpc_parse_connect_reply(res)
		print_status("#{rhost} Connected to #{self.wdbrpc_info[:rt_vers]} - #{self.wdbrpc_info[:rt_bsp_name]} (#{self.wdbrpc_info[:rt_bootline]})")

		report_note(
			:host   => rhost,
			:port   => rport,
			:proto  => 'udp',
			:type   => 'vxworks.target_info',
			:data   => res,
			:update => :unique
		)
	end

	def wdbrpc_client_memread(offset, length, params=0)
		pkt = wdbrpc_request_memread(offset, length, params)
		cnt = 0
		res = nil

		begin
			udp_sock.sendto(pkt, rhost, rport, 0)
			res,src = udp_sock.recvfrom(65535, 0.5)
			if not res and src
				raise RuntimeError, "no reply"
			end

			if res.length <= 48
				raise RuntimeError, "short read"
			end

		rescue ::Interrupt
			raise $!
		rescue ::Exception
			if cnt < 120
				cnt += 1
				retry
			end
		end

		res[48,res.length-48]
	end

	def wdbrpc_client_memwrite(offset, buffer, params=0)
		pkt = wdbrpc_request_memwrite(offset, buffer, params)
		cnt = 0
		res = nil

		udp_sock.sendto(pkt, rhost, rport, 0)
		res,src = udp_sock.recvfrom(65535, 5.0)

		if not res and src
			raise RuntimeError, "no reply"
		end
		res[-4,4].unpack("N")[0]
	end


	def wdbrpc_client_memscan(offset, depth, buffer, params=0)
		pkt = wdbrpc_request_memscan(offset, depth, buffer, params)
		cnt = 0
		res = nil

		udp_sock.sendto(pkt, rhost, rport, 0)
		res,src = udp_sock.recvfrom(65535, 5.0)

		if not res and src
			raise RuntimeError, "no reply"
		end
		p res
		res
	end


	def wdbrpc_client_context_kill(ctx_type=0, ctx=0)
		pkt = wdbrpc_request_context_kill(ctx_type, ctx)
		res = nil

		begin
			udp_sock.sendto(pkt, rhost, rport, 0)
			res,src = udp_sock.recvfrom(65535, 0.5)

		rescue ::Interrupt
			raise $!
		rescue ::Exception
		end
		res
	end

	def wdbrpc_client_send_disconnect
		pkt = wdbrpc_request_disconnect
		begin
			if self.udp_sock
				self.udp_sock.sendto(pkt, rhost, rport, 0)
				self.udp_sock.recvfrom(65535, 5)
			end
		rescue ::Interrupt
			raise $!
		rescue ::Exception
		end
	end

	def wdbrpc_client_disconnect
		wdbrpc_client_send_disconnect

		if self.udp_sock
			self.udp_sock.close rescue nil
		end
		self.udp_sock = nil

	end

	def rhost
		datastore['RHOST']
	end

	def rport
		datastore['RPORT'].to_i
	end

end
end

