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('ARP_SECRET', [true, 'The 32-bit cookie for ARP and UDP probe requests.', 1297303091]),
				OptAddress.new('GATEWAY', [false, 'The gateway IP address']),
			], Msf::Exploit::Capture
		)

		require 'racket'

		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={})

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

		# Capture device
		dev = nil
		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 = datastore['INTERFACE'] || ::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['ARP_SECRET'].to_i
				arp_filter = "arp host #{IPAddr.new(preamble, Socket::AF_INET).to_s} 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) 
			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 = Thread.new 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)
		if not pcap
			raise RuntimeError, "Could not access the capture process (remember to open_pcap first!)"
		else
			pcap.inject(pkt)
		end
	end

	# Injects an Ethernet packet with an optional payload.
	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
		n = Racket::Racket.new
		n.l2 = Racket::L2::Ethernet.new
		n.l2.dst_mac = eth_daddr
		n.l2.src_mac = eth_saddr
		n.l2.ethertype = eth_type
		pkt = n.pack
		pkt += payload if payload
		inject pkt,pcap
	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)
		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
		dst_mac,src_mac = lookup_eth(dhost)
		if dst_mac == "ff:ff:ff:ff:ff:ff" and not bcast
			return false
		end
		inject_eth(:payload => payload, :eth_daddr => dst_mac, :eth_saddr => src_mac)
	end

	# Depending on what kind of packet you get, the resultant hash returned will
	# contain one or several Racket objects.
	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|
						eth = Racket::L2::Ethernet.new(r)
						case proto
						when :arp
							next if not eth.ethertype == 0x0806
							arp = Racket::L3::ARP.new(eth.payload)
							reply = {:raw => r, :eth => eth, :arp => arp}
							break
						when :ip
							next if not eth.ethertype == 0x0800
							ip = Racket::L3::IPv4.new(eth.payload)
							reply = {:raw => r, :eth => eth, :ip => ip}
							break
						when :udp
							ip = Racket::L3::IPv4.new(eth.payload)
							next if not ip.protocol == 17
							udp = Racket::L4::UDP.new(ip.payload)
							reply = {:raw => r, :eth => eth, :ip => ip, :udp => udp, :payload => udp.payload}
							break
						when :tcp
							ip = Racket::L3::IPv4.new(eth.payload)
							next if not ip.protocol == 6
							tcp = Racket::L4::TCP.new(ip.payload)
							reply = {:raw => r, :eth => eth, :ip => ip, :tcp => tcp, :payload => tcp.payload}
							break
						end
					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
	#
	# This is done by first generating two packets: First, we try to ARP
	# the target IP address. If it's on the local network, this will usually
	# elicit a response of [dst_mac, dst_ip] -- see arp, below.
	#
	# We still need to know our own src_mac, though, since we assume we want
	# a response. To do this, we send a UDP trigger packet with a normal 
	# UDPSocket connect, to either the gateway host (GATEWAY), if known, or to 
	# a random(ish) IP known to be beyond the default gateway. We then listen 
	# for our own UDP packet (via inject_reply) -- it need not get a response 
	# from GATEWAY or the random fake IP.
	#
	# Note, if IANA assigns 177/8 in the future, then the default fake-remote
	# IP address will have to be changed (unless you don't mind sending them
	# probe packets).
	#
	# Finally, if all of this fails to get an adequate response, a default
	# ethernet address pair of ["ff:ff:ff:ff:ff:ff","00:00:00:00:00:00"] will
	# be used by the inject_eth function (above), unless the calling module
	# decides against it. In switched networks, you should always get a correct
	# pair for a remote address, the same pair for an unroutable address 
	# (eg, 127.0.0.1), or a correct pair for a local address (if the target 
	# responds to ARPs at all).
	#
	# TODO: Fix PcapRub to return a subnet mask so I can determine ahead of
	# time if it's in the local network.
	# 
	# TODO: Reorganize this. 
	#
	def lookup_eth(addr=nil,iface=nil)
		self.arp_cache = {} unless self.arp_cache.kind_of? Hash
		eth_pair = [nil,nil] # [dst_mac, src_mac] will go here.
		raise RuntimeError, "Could not access the capture process." if not self.arp_capture
		to = (datastore['TIMEOUT'] || 500).to_f / 1000.0
		# Try to ARP the real IP first.
		eth_pair[0] = self.arp_cache[addr] || arp(addr) || self.arp_cache[:gateway]
		eth_pair[1] = self.arp_cache[Rex::Socket.source_address(addr)]
		if (eth_pair[0].nil? || eth_pair[1].nil?)
			dst_host = (datastore['GATEWAY'] || IPAddr.new((rand(16777216) + 2969567232), Socket::AF_INET).to_s)
			dst_port = rand(30000)+1024
			preamble = [datastore['ARP_SECRET']].pack("N")
			secret = "#{preamble}#{Rex::Text.rand_text(rand(0xff)+1)}"
			UDPSocket.open.send(secret,0,dst_host,dst_port)
			begin
				Timeout.timeout(to) do
					while(my_packet = inject_reply(:udp,self.arp_capture))
						if my_packet[:payload] = secret
							eth_pair[0] ||= my_packet[:eth].dst_mac
							eth_pair[1] = my_packet[:eth].src_mac
							self.arp_cache[:gateway] = my_packet[:eth].dst_mac
							self.arp_cache[Rex::Socket.source_address(addr)] = my_packet[:eth].src_mac
							return eth_pair
						else
							next
						end
					end
				end
				rescue Timeout::Error
			end
		else
			return eth_pair
		end
	end

	# A pure-Ruby ARP exchange. It always uses self.arp_capture.
	def arp(target_ip=nil)
		return self.arp_cache[target_ip] if self.arp_cache[target_ip]
		to = (datastore['TIMEOUT'] || 500).to_f / 1000.0
		raise RuntimeError, "Could not access the capture process." if not self.arp_capture
		n = Racket::Racket.new
		n.l3 = Racket::L3::ARP.new
		n.l3.opcode = 1
		n.l3.tpa = target_ip || datastore['RHOST']
		n.l3.spa = IPAddr.new(datastore['ARP_SECRET'].to_i, Socket::AF_INET).to_s
		inject_eth(:eth_type => 0x0806,
							 :payload => n.pack,
							 :pcap => self.arp_capture
							)
		begin
			Timeout.timeout(to) do
				while (my_packet = inject_reply(:arp,self.arp_capture))
					if my_packet[:arp].spa == target_ip
						self.arp_cache[target_ip] = my_packet[:arp].sha
						return self.arp_cache[target_ip] 
					else
						next
					end
				end
			end
		rescue Timeout::Error
		end
	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
		
	attr_accessor :capture, :arp_cache, :arp_capture

end

end