# -*- 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('SECRET', [true, 'A 32-bit cookie for probe requests.', 'MSF!'.unpack('N').first]),
            OptAddress.new('GATEWAY_PROBE_HOST',
                           [
                             true,
                             'Send a TTL=1 random UDP datagram to this host to discover the default gateway\'s MAC',
                             'www.metasploit.com']),
            OptPort.new('GATEWAY_PROBE_PORT',
                        [
                          false,
                          'The port on GATEWAY_PROBE_HOST to send a random UDP probe to (random if 0 or unset)'])
          ], Msf::Exploit::Capture
        )

        require 'packetfu'

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

        begin
          require 'network_interface'
          @network_interface_loaded = true
        rescue ::LoadError => 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']
        do_arp = (opts['ARPCAP'] == false) ? false : 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 do_arp
            self.arp_capture = ::Pcap.open_live(dev, 512, true, tim)
            preamble         = datastore['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 (do_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 = {}
        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
      # to 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_PROBE_HOST']
        dst_port = datastore['GATEWAY_PROBE_PORT'] == 0 ? rand(30000) + 1024 : datastore['GATEWAY_PROBE_PORT']
        preamble = [datastore['SECRET']].pack("N")
        secret   = "#{preamble}#{Rex::Text.rand_text(rand(0xff)+1)}"

        begin
          UDPSocket.open do |sock|
            sock.setsockopt(::Socket::IPPROTO_IP, ::Socket::IP_TTL, 1)
            sock.send(secret, 0, dst_host, dst_port)
          end
        rescue Errno::ENETUNREACH
          # This happens on networks with no gateway. 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

        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
        begin
          my_ip, my_mask = Pcap.lookupnet(dev)
          # convert the netmask obtained from the relevant interface to CIDR
          cidr_mask = my_mask.to_s(2).count('1')
          my_net = IPAddr.new("#{my_ip}/#{cidr_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)
        lookupnet.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} does not exist" if !addrs
        raise RuntimeError, "Interface #{dev} does 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} does not exist" if !addrs
        raise RuntimeError, "Interface #{dev} does 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