#!/usr/bin/env ruby
require 'test/unit'
$:.unshift File.expand_path(File.join(File.dirname(__FILE__), "..", "lib"))
require 'packetfu'

class NewPacketTest < Test::Unit::TestCase
	include PacketFu
	
	def test_method_missing_and_respond_to
		p = TCPPacket.new
		assert p.respond_to?(:ip_len)
		assert p.ip_len = 20
		assert !(p.respond_to? :ip_bogus_header)
		assert_raise NoMethodError do
			p.bogus_header = 20
		end
	end

	def test_more_method_missing_magic
		p = UDPPacket.new
		assert_kind_of(UDPPacket,p)
		assert p.is_udp?
		assert p.is_ip?
		assert p.is_eth?
		assert_equal(p.ip_hl,5)
		assert p.layer
		assert_raise NoMethodError do 
			p.is_blue? 
		end
		assert_raise NoMethodError do
		 	p.tcp_blue
		end
		assert_raise NoMethodError do 
			p.udp_blue 
		end
		assert_raise NoMethodError do
			p.blue
		end
	end
end

class PacketStrippingTest < Test::Unit::TestCase

	include PacketFu

	def test_arp_strip
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[5], :fix => true) # Really ARP request.
		assert_kind_of(Packet,p)
		assert_kind_of(ARPPacket,p)
	end

end

class PacketParsersTest < Test::Unit::TestCase
	include PacketFu

	def test_parse_eth_packet
		assert_equal(EthPacket.layer, 1)
		assert_equal(EthPacket.layer_symbol, :link)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[5]) # Really ARP.
		assert_kind_of(Packet,p)
		assert_kind_of(EthHeader, p.headers[0])
		assert p.is_eth?
		assert_equal(pcaps[5],p.to_s)
	end

	def test_parse_arp_request
		assert_equal(ARPPacket.layer, 2)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[5]) # Really ARP request.
		assert p.is_eth?
		assert_kind_of(EthPacket,p)
		assert_kind_of(ARPPacket,p)
		assert p.is_arp?
		assert_equal(p.to_s, pcaps[5])
		assert_equal(1, p.arp_opcode.to_i)
		assert_equal("\x00\x01", p.headers.last[:arp_opcode].to_s)
	end

	def test_parse_arp_reply
		assert_equal(ARPPacket.layer, 2)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[6]) # Really ARP reply.
		assert_equal(p.to_s, pcaps[6])
		assert_equal(2, p.arp_opcode.to_i)
		assert_equal("\x00\x02", p.headers.last[:arp_opcode].to_s)
	end

	def test_parse_ip_packet
		assert_equal(IPPacket.layer, 2)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[0]) # Really DNS request 
		assert_equal(p.to_s[0,20], pcaps[0][0,20])
		assert_equal(p.to_s, pcaps[0])
		assert_kind_of(EthPacket,p)
		assert_kind_of(IPPacket,p)
	end

	def test_parse_tcp_packet
		assert_equal(TCPPacket.layer, 3)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[7]) # Really FIN/ACK
		assert_equal(p.to_s, pcaps[7])
		assert_kind_of(EthPacket,p)
		assert_kind_of(IPPacket,p)
		assert_kind_of(TCPPacket,p)
	end

	def test_parse_udp_packet
		assert_equal(UDPPacket.layer, 3)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[0]) # Really DNS request
		assert_equal(p.to_s, pcaps[0])
		assert_kind_of(EthPacket,p)
		assert_kind_of(IPPacket,p)
		assert_kind_of(UDPPacket,p)
	end

	def test_parse_icmp_packet
		assert_equal(ICMPPacket.layer, 3)
		assert_equal(ICMPPacket.layer_symbol, :transport)
		pcaps = PcapFile.new.file_to_array(:f => 'sample.pcap')
		p = Packet.parse(pcaps[3]) # Really ICMP reply
		assert_equal(p.to_s, pcaps[3])
		assert_kind_of(EthPacket,p)
		assert_kind_of(IPPacket,p)
		assert_kind_of(ICMPPacket,p)
	end

	def test_parse_invalid_packet
		assert_equal(InvalidPacket.layer, 0)
		assert_equal(InvalidPacket.layer_symbol, :invalid)
		p = Packet.parse("\xff\xfe\x00\x01")
		assert_equal(p.to_s, "\xff\xfe\x00\x01")
		assert_kind_of(InvalidPacket,p)
	end

	def test_parse_ipv6_packet
		assert_equal(IPv6Packet.layer, 2)
		assert_equal(IPv6Packet.layer_symbol, :internet)
		pcaps = PcapFile.new.file_to_array(:f => 'sample-ipv6.pcap')
		p = Packet.parse(pcaps[0]) # Really an IPv6 packet 
		assert_equal(p.to_s, pcaps[0])
		assert_kind_of(EthPacket,p)
		assert(!p.kind_of?(IPPacket), "Misidentified as an IP Packet!")
		assert_kind_of(IPv6Packet,p)
	end

	def test_parse_hsrp_packet
		assert_equal(HSRPPacket.layer, 4)
		assert_equal(HSRPPacket.layer_symbol, :application)
		pcaps = PcapFile.new.file_to_array(:f => 'sample_hsrp_pcapr.cap')
		p = Packet.parse(pcaps[0]) # Really an HSRP Hello packet 
		assert_equal(p.to_s, pcaps[0])
		assert_kind_of(EthPacket,p)
		assert_kind_of(IPPacket,p)
		assert_kind_of(UDPPacket,p)
		assert_kind_of(HSRPPacket,p)
	end

	def test_parse_hsrp_as_udp
		assert_equal(:application, HSRPPacket.layer_symbol)
		pcaps = PcapFile.new.file_to_array(:f => 'sample_hsrp_pcapr.cap')
		p = Packet.parse(pcaps[0], :parse_app => false) # Really an HSRP Hello packet 
		assert_kind_of(UDPPacket,p)
		assert(!p.kind_of?(HSRPPacket), "Misidentified HSRP packet when we didn't want it!" )
	end

end


# vim: nowrap sw=2 sts=0 ts=2 ff=unix ft=ruby
