# -*- coding: binary -*-
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.
#
###

	class Exploit
		module 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

				begin
					require 'network_interface'
					@network_interface_loaded = true
				rescue ::Exception => e
					@network_interface_loaded = false
					@network_interface_error  = e
				end

			end

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

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

			def stats_ifdrop(pcap=self.capture)
				return(0) unless 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

					unless RUBY_PLATFORM == "i386-mingw32"
						system("ifconfig", dev, "up")
					end

					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 unless 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 unless 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 unless 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
				unless 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)
					Rex.sleep((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." unless 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." unless 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"
				elsif not @network_interface_loaded
					print_status("The NetworkInterface module is not available: #{@network_interface_error}")
					raise RuntimeError, "NetworkInterface 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

			def netifaces_implemented?
				@network_interface_loaded and
					NetworkInterface.respond_to?(:interfaces) and
					NetworkInterface.respond_to?(:addresses)
			end

			def list_interfaces
				check_pcaprub_loaded
				NetworkInterface.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 NetworkInterface.interfaces.include?(dev)
					elsif dev.to_s =~ /^[0-9]{1,2}$/
						if (dev.to_i <= NetworkInterface.interfaces.length) and (dev.to_i >= 0)
							return true
						else
							return false
						end
					else
						return false
					end
				else
					return NetworkInterface.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)
							NetworkInterface.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 = NetworkInterface.addresses(dev)
				raise RuntimeError, "Interface #{dev} does not exist" if !addrs
				raise RuntimeError, "Can not get mac address for interface #{dev}" if !addrs[NetworkInterface::AF_LINK][0]['addr']
				addrs[NetworkInterface::AF_LINK][0]['addr']
			end

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

			def get_ipv4_addr(dev, num=0)
				check_pcaprub_loaded
				dev   = get_interface_guid(dev)
				addrs = NetworkInterface.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[NetworkInterface::AF_INET].length < num + 1
				raise RuntimeError, "Can not get the IPv4 address for interface #{dev}" if !addrs[NetworkInterface::AF_INET][num]['addr']
				addrs[NetworkInterface::AF_INET][num]['addr']
			end

			def get_ipv4_netmask(dev, num=0)
				check_pcaprub_loaded
				dev   = get_interface_guid(dev)
				addrs = NetworkInterface.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[NetworkInterface::AF_INET].length < num + 1
				raise RuntimeError, "Can not get IPv4 netmask for interface #{dev}" if !addrs[NetworkInterface::AF_INET][num]['netmask']
				addrs[NetworkInterface::AF_INET][num]['netmask']
			end

			def get_ipv4_broadcast(dev, num=0)
				check_pcaprub_loaded
				dev   = get_interface_guid(dev)
				addrs = NetworkInterface.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[NetworkInterface::AF_INET].length < num + 1
				raise RuntimeError, "Can not get IPv4 broadcast address for interface #{dev}" if !addrs[NetworkInterface::AF_INET][num]['broadcast']
				addrs[NetworkInterface::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" unless ::NetworkInterface.const_defined?(:AF_INET6)
				addrs = NetworkInterface.addresses(dev)
				raise RuntimeError, "Interface #{dev} do not exists" if !addrs
				addrs[NetworkInterface::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" unless ::NetworkInterface.const_defined?(:AF_INET6)
				addrs = NetworkInterface.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[NetworkInterface::AF_INET6].length < num + 1
				raise RuntimeError, "Can not get ipv6 address for interface #{dev}" if !addrs[NetworkInterface::AF_INET6][num]['addr']
				addrs[NetworkInterface::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" unless ::NetworkInterface.const_defined?(:AF_INET6)
				addrs = NetworkInterface.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[NetworkInterface::AF_INET6].length < num + 1
				raise RuntimeError, "Can not get ipv6 netmask address for interface #{dev}" if !addrs[NetworkInterface::AF_INET6][num]['netmask']
				addrs[NetworkInterface::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

end
