module Msf

###
#
# This module provides methods for sending and receiving
# raw packets. It should be preferred over the soon-to-be
# deprecated Rex::Socket::Ip and Msf::Exploite::Remote::Ip
# mixins.
#
# Please see the pcaprub documentation for more information
# on how to use capture objects.
#
###

module Exploit::Capture

	#
	# Initializes an instance of an exploit module that captures traffic
	#

	def initialize(info = {})
		super

		register_options(
			[
				OptPath.new('PCAPFILE', [false, 'The name of the PCAP capture file to process']),
				OptString.new('INTERFACE', [false, 'The name of the interface']),
				OptString.new('FILTER', [false, 'The filter string for capturing traffic']),
				OptInt.new('SNAPLEN', [true, 'The number of bytes to capture', 65535]),
				OptInt.new('TIMEOUT', [true, 'The number of seconds to wait for new data', 500]),
				Opt::RHOST

			], Msf::Exploit::Capture
		)

		register_advanced_options(
			[
				OptInt.new('UDP_SECRET', [true, 'The 32-bit cookie for UDP probe requests.', 1297303091]),
				OptAddress.new('GATEWAY', [false, 'The gateway IP address. This will be used rather than a random remote address for the UDP probe, if set.']),
				OptInt.new('NETMASK', [false, 'The local network mask. This is used to decide if an address is in the local network.', 24]),
			], Msf::Exploit::Capture
		)

		require 'packetfu'

		begin
			require 'pcaprub'
			@pcaprub_loaded = true
		rescue ::Exception => e
			@pcaprub_loaded = false
			@pcaprub_error  = e
		end

	end

	def stats_recv(pcap=self.capture)
		return(0) if not pcap
		pcap.stats['recv']
	end

	def stats_drop(pcap=self.capture)
		return(0) if not pcap
		pcap.stats['drop']
	end

	def stats_ifdrop(pcap=self.capture)
		return(0) if not pcap
		pcap.stats['ifdrop']
	end

	#
	# Opens a handle to the specified device
	#
	def open_pcap(opts={})
		check_pcaprub_loaded
		if RUBY_PLATFORM == "i386-mingw32"
			if opts['INTERFACE'] or datastore['INTERFACE']
				dev = opts['INTERFACE'] || datastore['INTERFACE']
				if is_interface?(dev)
					dev = get_interface_guid(dev)
				end
			end
		else
			dev = opts['INTERFACE'] || datastore['INTERFACE'] || nil
		end

		len = (opts['SNAPLEN'] || datastore['SNAPLEN']  || 65535).to_i
		tim = (opts['TIMEOUT'] || datastore['TIMEOUT']  || 0).to_i
		fil = opts['FILTER'] || datastore['FILTER']
		arp = opts['ARPCAP'] || true

		# Look for a PCAP file
		cap = datastore['PCAPFILE'] || ''

		if(not cap.empty?)
			if(not File.exists?(cap))
				raise RuntimeError, "The PCAP file #{cap} could not be found"
			end
			self.capture = ::Pcap.open_offline(cap)
		else
			dev ||= ::Pcap.lookupdev
			system("ifconfig", dev, "up")
			self.capture = ::Pcap.open_live(dev, len, true, tim)
			if arp
				self.arp_capture = ::Pcap.open_live(dev, 512, true, tim)
				preamble = datastore['UDP_SECRET'].to_i
				arp_filter = "arp[6:2] = 2 or (udp[8:4] = #{preamble})"
				self.arp_capture.setfilter(arp_filter)
			end
		end

		if (not self.capture)
			raise RuntimeError, "Could not start the capture process"
		elsif (arp and !self.arp_capture and cap.empty?)
			raise RuntimeError, "Could not start the ARP capture process"
		end

		self.capture.setfilter(fil) if fil
	end

	def close_pcap
		return if not self.capture
		self.capture = nil
		self.arp_capture = nil
		GC.start()
	end

	def capture_extract_ies(raw)
		set = {}
		ret = 0
		idx = 0
		len = 0

		while (idx < raw.length)
			len = raw[idx+1]
			return set if not len
			set[ raw[idx] ] ||= []
			set[ raw[idx] ].push(raw[idx + 2, len])
			idx += len + 2
		end

		return set
	end

	#
	# This monstrosity works around a series of bugs in the interrupt
	# signal handling of Ruby 1.9
	#
	def each_packet
		return if not capture
		begin
			@capture_count = 0
			reader = framework.threads.spawn("PcapReceiver", false) do
				capture.each do |pkt|
					yield(pkt)
					@capture_count += 1
				end
			end
			reader.join
		rescue ::Exception
			raise $!
		ensure
			reader.kill if reader.alive?
		end

		@capture_count
	end

	# Injects a packet on the wire. For all injection-related functions, it's
	# on the module to open up a capture device first (this way, we don't
	# needlessly spawn new capture devices).
	def inject(pkt="",pcap=self.capture)
		check_pcaprub_loaded
		if not pcap
			raise RuntimeError, "Could not access the capture process (remember to open_pcap first!)"
		else
			pcap.inject(pkt.to_s) # Can be a PacketFu Packet object or a pre-packed string
		end
	end

	# Injects an Ethernet packet with an optional payload. The payload
	# may be a regular PacketFu packet, an EthHeader, or a string.
	def inject_eth(args={})
		eth_daddr = args[:eth_daddr] || "ff:ff:ff:ff:ff:ff"
		eth_saddr = args[:eth_saddr] || "00:00:00:00:00:00"
		eth_type  = args[:eth_type]  || 0x0800 # IP default
		payload   = args[:payload]
		pcap      = args[:pcap]      || self.capture
		p = PacketFu::EthPacket.new
		p.eth_daddr = eth_daddr
		p.eth_saddr = eth_saddr
		p.eth_proto = eth_type
		if payload
			if payload.kind_of? PacketFu::EthPacket
				p.payload = payload.eth_header.body
			elsif payload.kind_of? PacketFu::EthHeader
				p.payload = payload.body
			else
				p.payload = payload.to_s
			end
		end
		inject p.to_s,pcap
	end

	def inject_pcap(pcap_file, filter=nil, delay = 0, pcap=self.capture)
		check_pcaprub_loaded
		if not pcap
			raise RuntimeError, "Could not access the capture process (remember to open_pcap first!)"
		end

		if(not File.exists?(pcap_file))
			raise RuntimeError, "The PCAP file #{pcap_file} could not be found"
		end

		if(pcap_file.empty?)
			raise RuntimeError, "The PCAP file #{pcap_file} is empty"
		end

		capture_file = ::Pcap.open_offline(pcap_file)
		capture_file.setfilter(filter) if filter
		while (pkt = capture_file.next) do
			pcap.inject(pkt)
			Kernel.select(nil, nil, nil, (delay * 1.0)/1000)
		end
		GC.start
	end

	# Capture_sendto is intended to replace the old Rex::Socket::Ip.sendto method. It requires
	# a payload and a destination address. To send to the broadcast address, set bcast
	# to true (this will guarantee that packets will be sent even if ARP doesn't work
	# out).
	def capture_sendto(payload="", dhost=nil, bcast=false, dev=nil)
		raise RuntimeError, "Could not access the capture process (remember to open_pcap first!)" unless self.capture
		raise RuntimeError, "Must specify a host to sendto" unless dhost
		dev ||= datastore['INTERFACE']
		dst_mac,src_mac = lookup_eth(dhost,dev)
		if dst_mac == nil and not bcast
			return false
		end
		inject_eth(:payload => payload, :eth_daddr => dst_mac, :eth_saddr => src_mac)
	end

	# The return value either be a PacketFu::Packet object, or nil
	def inject_reply(proto=:udp,pcap=self.capture)
		reply = nil
		to = (datastore['TIMEOUT'] || 500).to_f / 1000.0
		if not pcap
			raise RuntimeError, "Could not access the capture process (remember to open_pcap first!)"
		else
			begin
				::Timeout.timeout(to) do
					pcap.each do |r|
						packet = PacketFu::Packet.parse(r)
						next unless packet.proto.map {|x| x.downcase.to_sym}.include? proto
						reply = packet
						break
					end
				end
				rescue ::Timeout::Error
			end
		end
		return reply
	end

	# This ascertains the correct Ethernet addresses one should use to
	# ensure injected IP packets actually get where they are going, and
	# manages the self.arp_cache hash. It always uses self.arp_capture
	# do inject and capture packets, and will always first fire off a
	# UDP packet using the regular socket to learn the source host's
	# and gateway's mac addresses.
	def lookup_eth(addr=nil,iface=nil)
		raise RuntimeError, "Could not access the capture process." if not self.arp_capture

		self.arp_cache ||= {}
		self.dst_cache ||= {}

		return self.dst_cache[addr] if self.dst_cache[addr]

		if !self.arp_cache[Rex::Socket.source_address(addr)]
			probe_gateway(addr)
		end

		src_mac = self.arp_cache[Rex::Socket.source_address(addr)]
		if should_arp?(addr)
			dst_mac = self.arp_cache[addr] || arp(addr)
		else
			dst_mac = self.arp_cache[:gateway]
		end

		self.dst_cache[addr] = [dst_mac,src_mac]
	end

	def probe_gateway(addr)
		dst_host = (datastore['GATEWAY'] || IPAddr.new((rand(16777216) + 2969567232), Socket::AF_INET).to_s)
		dst_port = rand(30000)+1024
		preamble = [datastore['UDP_SECRET']].pack("N")
		secret = "#{preamble}#{Rex::Text.rand_text(rand(0xff)+1)}"
		UDPSocket.open.send(secret,0,dst_host,dst_port)
		begin
			to = (datastore['TIMEOUT'] || 1500).to_f / 1000.0
			::Timeout.timeout(to) do
				while(my_packet = inject_reply(:udp,self.arp_capture))
					if my_packet.payload == secret
						dst_mac = self.arp_cache[:gateway] = my_packet.eth_daddr
						src_mac = self.arp_cache[Rex::Socket.source_address(addr)] = my_packet.eth_saddr
						return [dst_mac,src_mac]
					else
						next
					end
				end
			end
		rescue ::Timeout::Error
			# Well, that didn't work (this common on networks where there's no gatway, like
			# VMWare network interfaces. We'll need to use a fake source hardware address.
			self.arp_cache[Rex::Socket.source_address(addr)] = "00:00:00:00:00:00"
		end
	end

	# A pure-Ruby ARP exchange. It uses self.arp_capture to send and recv
	# packets, rather than self.capture.
	def arp(target_ip=nil)
		return self.arp_cache[target_ip] if self.arp_cache[target_ip]
		return self.arp_cache[:gateway] unless should_arp? target_ip
		source_ip = Rex::Socket.source_address(target_ip)
		raise RuntimeError, "Could not access the capture process." if not self.arp_capture
		p = arp_packet(target_ip,source_ip)
		inject_eth(:eth_type => 0x0806,
			:payload => p,
			:pcap => self.arp_capture,
			:eth_saddr => self.arp_cache[Rex::Socket.source_address(target_ip)]
		)
		begin
			to = (datastore['TIMEOUT'] || 500).to_f / 1000.0
			::Timeout.timeout(to) do
				while(my_packet = inject_reply(:arp,self.arp_capture))
					if my_packet.arp_saddr_ip == target_ip
						self.arp_cache[target_ip] = my_packet.eth_saddr
						return self.arp_cache[target_ip]
					else
						next
					end
				end
			end
		rescue ::Timeout::Error
		end
	end

	# Creates a full ARP packet, mainly for use with inject_eth()
	def arp_packet(target_ip=nil,source_ip=nil)
		p = PacketFu::ARPPacket.new
		p.arp_opcode = 1
		p.arp_daddr_ip = target_ip || datastore['RHOST']
		p.arp_saddr_ip = source_ip || datastore['LHOST']
		my_eth = self.arp_cache[Rex::Socket.source_address(target_ip)]
		p.arp_saddr_mac = my_eth || "00:00:00:00:00:00"
		return p
	end

	# Allow modules to reset their arp caches arbitrarily.
	def expire_arpcache
		self.arp_cache = {}
	end

	# For compatabilty with Msf::Exploit::Remote::Ip
	def rhost
		datastore['RHOST']
	end

	def check_pcaprub_loaded
		if not @pcaprub_loaded
			print_status("The Pcaprub module is not available: #{@pcaprub_error}")
			raise RuntimeError, "Pcaprub not available"
		else
			true
		end
	end

	def lookupnet
		check_pcaprub_loaded
		dev  = datastore['INTERFACE'] || ::Pcap.lookupdev
		mask = datastore['NETMASK'] || 24
		begin
			my_net = IPAddr.new("#{Pcap.lookupnet(dev).first}/#{mask}")
		rescue RuntimeError => e
			@pcaprub_error = e
			print_status("Cannot stat device: #{@pcaprub_error}")
			raise RuntimeError, "Pcaprub error: #{@pcaprub_error}"
		end
		return my_net
	end

	def should_arp?(ip)
		@mydev  ||= datastore['INTERFACE'] || ::Pcap.lookupdev
		@mymask ||= datastore['NETMASK'] || 24
		@mynet  ||= lookupnet
		@mynet.include?(IPAddr.new(ip))
	end

	attr_accessor :capture, :arp_cache, :arp_capture, :dst_cache

	#Netifaces code

	# netifaces code is  not available in pcaprub 0.9.2 and prior,
	# which is going to be installed in a lot of places. Modules
	# which want it should check explicitly for it. TODO: Bug upstream
	# to release it for real in 0.9.3
	def netifaces_implemented?
		@pcaprub_loaded and
		Pcap.respond_to?(:lookupaddrs) and
		Pcap.respond_to?(:interfaces) and
		Pcap.respond_to?(:addresses)
	end

	def list_interfaces
		check_pcaprub_loaded
		Pcap.interfaces
	end

	def is_interface?(dev)
		check_pcaprub_loaded
		if RUBY_PLATFORM == "i386-mingw32"
			if dev =~ /\\Device\\NPF_\{[A-F0-9]{8}-[A-F0-9]{4}-[A-F0-9]{4}-[A-F0-9]{4}-[A-F0-9]{12}\}/
				return Pcap.interfaces.include?(dev)
			elsif dev.to_s =~ /^[0-9]{1,2}$/
				if (dev.to_i <= Pcap.interfaces.length) and (dev.to_i >= 0)
					return true
				else
					return false
				end
			else
				return false
			end
		else
			return Pcap.interfaces.include?(dev)
		end
	end

	# This function is usefull only on windows where pcaprub use the GUID
	def get_interface_guid(dev)
		check_pcaprub_loaded
		if RUBY_PLATFORM == "i386-mingw32"
			if dev.to_s =~ /^[0-9]{1,2}$/
				if is_interface?(dev)
					Pcap.interfaces[(dev.to_i) - 1]
				else
					return dev
				end
			else
				return dev
			end
		else #Non windows
			return dev
		end
	end

	def get_mac(dev)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} does not exist" if !addrs
		raise RuntimeError, "Can not get mac address for interface #{dev}" if !addrs[Pcap::AF_LINK][0]['addr']
		addrs[Pcap::AF_LINK][0]['addr']
	end

	def get_ipv4_addr_count(dev)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} does not exist" if !addrs
		addrs[Pcap::AF_INET].length
	end

	def get_ipv4_addr(dev, num=0)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} do not exists" if !addrs
		raise RuntimeError, "Interface #{dev} do not have an ipv4 address at position #{num}" if addrs[Pcap::AF_INET].length < num + 1
		raise RuntimeError, "Can not get the IPv4 address for interface #{dev}" if !addrs[Pcap::AF_INET][num]['addr']
		addrs[Pcap::AF_INET][num]['addr']
	end

	def get_ipv4_netmask(dev, num=0)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} do not exists" if !addrs
		raise RuntimeError, "Interface #{dev} do not have an ipv4 address at position #{num}" if addrs[Pcap::AF_INET].length < num + 1
		raise RuntimeError, "Can not get IPv4 netmask for interface #{dev}" if !addrs[Pcap::AF_INET][num]['netmask']
		addrs[Pcap::AF_INET][num]['netmask']
	end

	def get_ipv4_broadcast(dev, num=0)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} do not exists" if !addrs
		raise RuntimeError, "Interface #{dev} do not have an ipv4 address at position #{num}" if addrs[Pcap::AF_INET].length < num + 1
		raise RuntimeError, "Can not get IPv4 broadcast address for interface #{dev}" if !addrs[Pcap::AF_INET][num]['broadcast']
		addrs[Pcap::AF_INET][num]['broadcast']
	end

	def get_ipv6_addr_count(dev)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		raise RuntimeError, "IPv6 information is not available on this platform" if not ::Pcap.const_defined?(:AF_INET6)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} do not exists" if !addrs
		addrs[Pcap::AF_INET6].length
	end

	# NOTE: IPv6 is not implemented on Windows
	def get_ipv6_addr(dev, num=0)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		raise RuntimeError, "IPv6 information is not available on this platform" if not ::Pcap.const_defined?(:AF_INET6)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} do not exists" if !addrs
		raise RuntimeError, "Interface #{dev} do not have an ipv6 address at position #{num}" if addrs[Pcap::AF_INET6].length < num + 1
		raise RuntimeError, "Can not get ipv6 address for interface #{dev}" if !addrs[Pcap::AF_INET6][num]['addr']
		addrs[Pcap::AF_INET6][num]['addr'].gsub(/%(.)*$/,'')
	end

	def get_ipv6_netmask(dev, num=0)
		check_pcaprub_loaded
		dev = get_interface_guid(dev)
		raise RuntimeError, "IPv6 information is not available on this platform" if not ::Pcap.const_defined?(:AF_INET6)
		addrs = Pcap.addresses(dev)
		raise RuntimeError, "Interface #{dev} do not exists" if !addrs
		raise RuntimeError, "Interface #{dev} do not have an ipv6 address at position #{num}" if addrs[Pcap::AF_INET6].length < num + 1
		raise RuntimeError, "Can not get ipv6 netmask address for interface #{dev}" if !addrs[Pcap::AF_INET6][num]['netmask']
		addrs[Pcap::AF_INET6][num]['netmask']
	end

	# Protocol-specific encoding/decoding methods until more
	# application protos get into PacketFu proper

	# Intended to be used as the payload to an ICMP echo request's payload
	def capture_icmp_echo_pack(id=nil,seq=nil,payload=nil)
		id ||= rand(0x10000)
		seq ||= rand(0x10000)
		[id,seq,payload.to_s].pack("nna*")
	end

	# Decodes and ICMP echo request or response.
	def capture_icmp_echo_unpack(data)
		data.unpack("nna*")
	end

end

end