1 files changed, 80 insertions, 107 deletions
diff --git a/cucumber/features/support/helpers/firewall_helper.rb b/cucumber/features/support/helpers/firewall_helper.rb
index fce363c5..f88091de 100644
--- a/cucumber/features/support/helpers/firewall_helper.rb
+++ b/cucumber/features/support/helpers/firewall_helper.rb
@@ -1,121 +1,94 @@
 require 'packetfu'
-require 'ipaddr'
 
-# Extent IPAddr with a private/public address space checks
-class IPAddr
-  PrivateIPv4Ranges = [
-    IPAddr.new("10.0.0.0/8"),
-    IPAddr.new("172.16.0.0/12"),
-    IPAddr.new("192.168.0.0/16"),
-    IPAddr.new("255.255.255.255/32")
-  ]
-
-  PrivateIPv6Ranges = [
-    IPAddr.new("fc00::/7")
-  ]
-
-  def private?
-    private_ranges = self.ipv4? ? PrivateIPv4Ranges : PrivateIPv6Ranges
-    private_ranges.any? { |range| range.include?(self) }
-  end
-
-  def public?
-    !private?
-  end
+def looks_like_dhcp_packet?(eth_packet, protocol, sport, dport, ip_packet)
+  protocol == "udp" && sport == 68 && dport == 67 && 
+    eth_packet.eth_daddr == "ff:ff:ff:ff:ff:ff" &&
+    ip_packet && ip_packet.ip_daddr == "255.255.255.255"
 end
 
-class FirewallLeakCheck
-  attr_reader :ipv4_tcp_leaks, :ipv4_nontcp_leaks, :ipv6_leaks, :nonip_leaks, :mac_leaks
-
-  def initialize(pcap_file, options = {})
-    options[:accepted_hosts] ||= []
-    options[:ignore_lan] ||= true
-    @pcap_file = pcap_file
-    packets = PacketFu::PcapFile.new.file_to_array(:filename => @pcap_file)
-    mac_leaks = Set.new
-    ipv4_tcp_packets = []
-    ipv4_nontcp_packets = []
-    ipv6_packets = []
-    nonip_packets = []
-    packets.each do |p|
-      if PacketFu::EthPacket.can_parse?(p)
-        packet = PacketFu::EthPacket.parse(p)
-        mac_leaks << packet.eth_saddr
-        mac_leaks << packet.eth_daddr
-      end
-
-      if PacketFu::TCPPacket.can_parse?(p)
-        ipv4_tcp_packets << PacketFu::TCPPacket.parse(p)
-      elsif PacketFu::IPPacket.can_parse?(p)
-        ipv4_nontcp_packets << PacketFu::IPPacket.parse(p)
-      elsif PacketFu::IPv6Packet.can_parse?(p)
-        ipv6_packets << PacketFu::IPv6Packet.parse(p)
-      elsif PacketFu::Packet.can_parse?(p)
-        nonip_packets << PacketFu::Packet.parse(p)
-      else
-        save_pcap_file
-        raise "Found something in the pcap file that cannot be parsed"
-      end
+# Returns the unique edges (based on protocol, source/destination
+# address/port) in the graph of all network flows.
+def pcap_connections_helper(pcap_file, opts = {})
+  opts[:ignore_dhcp] = true unless opts.has_key?(:ignore_dhcp)
+  connections = Array.new
+  packets = PacketFu::PcapFile.new.file_to_array(:filename => pcap_file)
+  packets.each do |p|
+    if PacketFu::EthPacket.can_parse?(p)
+      eth_packet = PacketFu::EthPacket.parse(p)
+    else
+      raise 'Found something that is not an ethernet packet'
+    end
+    sport = nil
+    dport = nil
+    if PacketFu::IPv6Packet.can_parse?(p)
+      ip_packet = PacketFu::IPv6Packet.parse(p)
+      protocol = 'ipv6'
+    elsif PacketFu::TCPPacket.can_parse?(p)
+      ip_packet = PacketFu::TCPPacket.parse(p)
+      protocol = 'tcp'
+      sport = ip_packet.tcp_sport
+      dport = ip_packet.tcp_dport
+    elsif PacketFu::UDPPacket.can_parse?(p)
+      ip_packet = PacketFu::UDPPacket.parse(p)
+      protocol = 'udp'
+      sport = ip_packet.udp_sport
+      dport = ip_packet.udp_dport
+    elsif PacketFu::ICMPPacket.can_parse?(p)
+      ip_packet = PacketFu::ICMPPacket.parse(p)
+      protocol = 'icmp'
+    elsif PacketFu::IPPacket.can_parse?(p)
+      ip_packet = PacketFu::IPPacket.parse(p)
+      protocol = 'ip'
+    else
+      raise "Found something that cannot be parsed"
     end
-    ipv4_tcp_hosts = filter_hosts_from_ippackets(ipv4_tcp_packets,
-                                                 options[:ignore_lan])
-    accepted = Set.new(options[:accepted_hosts])
-    @mac_leaks = mac_leaks
-    @ipv4_tcp_leaks = ipv4_tcp_hosts.select { |host| !accepted.member?(host) }
-    @ipv4_nontcp_leaks = filter_hosts_from_ippackets(ipv4_nontcp_packets,
-                                                     options[:ignore_lan])
-    @ipv6_leaks = filter_hosts_from_ippackets(ipv6_packets,
-                                              options[:ignore_lan])
-    @nonip_leaks = nonip_packets
-  end
 
-  def save_pcap_file
-    save_failure_artifact("Network capture", @pcap_file)
-  end
+    next if opts[:ignore_dhcp] &&
+            looks_like_dhcp_packet?(eth_packet, protocol,
+                                    sport, dport, ip_packet)
 
-  # Returns a list of all unique destination IP addresses found in
-  # `packets`. Exclude LAN hosts if ignore_lan is set.
-  def filter_hosts_from_ippackets(packets, ignore_lan)
-    hosts = []
-    packets.each do |p|
-      candidate = nil
-      if p.kind_of?(PacketFu::IPPacket)
-        candidate = p.ip_daddr
-      elsif p.kind_of?(PacketFu::IPv6Packet)
-        candidate = p.ipv6_header.ipv6_daddr
-      else
-        save_pcap_file
-        raise "Expected an IP{v4,v6} packet, but got something else:\n" +
-              p.peek_format
-      end
-      if candidate != nil and (not(ignore_lan) or IPAddr.new(candidate).public?)
-        hosts << candidate
+    packet_info = {
+      mac_saddr: eth_packet.eth_saddr,
+      mac_daddr: eth_packet.eth_daddr,
+      protocol: protocol,
+      sport: sport,
+      dport: dport,
+    }
+
+    begin
+      packet_info[:saddr] = ip_packet.ip_saddr
+      packet_info[:daddr] = ip_packet.ip_daddr
+    rescue NoMethodError, NameError
+      begin
+        packet_info[:saddr] = ip_packet.ipv6_saddr
+        packet_info[:daddr] = ip_packet.ipv6_daddr
+      rescue NoMethodError, NameError
+        puts "We were hit by #11508. PacketFu bug? Packet info: #{ip_packet}"
+        packet_info[:saddr] = nil
+        packet_info[:daddr] = nil
       end
     end
-    hosts.uniq
+    connections << packet_info
   end
+  connections.uniq.map { |p| OpenStruct.new(p) }
+end
 
-  def assert_no_leaks
-    err = ""
-    if !@ipv4_tcp_leaks.empty?
-      err += "The following IPv4 TCP non-Tor Internet hosts were " +
-        "contacted:\n" + ipv4_tcp_leaks.join("\n")
-    end
-    if !@ipv4_nontcp_leaks.empty?
-      err += "The following IPv4 non-TCP Internet hosts were contacted:\n" +
-        ipv4_nontcp_leaks.join("\n")
-    end
-    if !@ipv6_leaks.empty?
-      err += "The following IPv6 Internet hosts were contacted:\n" +
-        ipv6_leaks.join("\n")
-    end
-    if !@nonip_leaks.empty?
-      err += "Some non-IP packets were sent\n"
-    end
-    if !err.empty?
-      save_pcap_file
-      raise err
-    end
+class FirewallAssertionFailedError < Test::Unit::AssertionFailedError
+end
+
+# These assertions are made from the perspective of the system under
+# testing when it comes to the concepts of "source" and "destination".
+def assert_all_connections(pcap_file, opts = {}, &block)
+  all = pcap_connections_helper(pcap_file, opts)
+  good = all.find_all(&block)
+  bad = all - good
+  unless bad.empty?
+    raise FirewallAssertionFailedError.new(
+            "Unexpected connections were made:\n" +
+            bad.map { |e| "  #{e}" } .join("\n"))
   end
+end
 
+def assert_no_connections(pcap_file, opts = {}, &block)
+  assert_all_connections(pcap_file, opts) { |*args| not(block.call(*args)) }
 end