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

class String
  def bin
    self.scan(/../).map {|x| x.to_i(16).chr}.join
  end
end

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

  def test_udp_header_new
    u = UDPHeader.new
    assert_kind_of UDPHeader, u
    assert_equal(8, u.to_s.size)
    assert_equal("\x00\x00\x00\x00\x00\x08\x00\x00", u.to_s)
  end

  def test_udp_peek
    u = UDPPacket.new
    u.ip_saddr = "10.20.30.40"
    u.ip_daddr = "50.60.70.80"
    u.udp_src = 53
    u.udp_dport = 1305
    u.payload = "abcdefghijklmnopqrstuvwxyz"
    u.recalc
    puts "\n"
    puts "UDP Peek format: "
    puts u.peek
    assert (u.peek.size <= 80)
  end

  def test_udp_pcap
    u = UDPPacket.new
    assert_kind_of UDPPacket, u
    u.recalc
    u.to_f('udp_test.pcap','a')
    u.ip_saddr = "10.20.30.40"
    u.ip_daddr = "50.60.70.80"
    u.payload = "+some fakey-fake udp packet"
    u.udp_src = 1205
    u.udp_dst = 13013
    u.recalc
    u.to_f('udp_test.pcap','a')
  end

  def test_udp_read
    sample_packet = PcapFile.new.file_to_array(:f => 'sample.pcap')[0]
    pkt = Packet.parse(sample_packet)
    assert_kind_of UDPPacket, pkt
    assert_equal(0x8bf8, pkt.udp_sum.to_i)
    pkt.to_f('udp_test.pcap','a')
  end

  def test_udp_checksum
    sample_packet = PcapFile.new.file_to_array(:f => 'sample.pcap')[0]
    pkt = Packet.parse(sample_packet)
    assert_kind_of UDPPacket, pkt
    pkt.recalc
    assert_equal(0x8bf8, pkt.udp_sum.to_i)
    pkt.to_f('udp_test.pcap','a')
  end

  def test_udp_read_strip
    str = "01005e7ffffa100ba9eb63400800450000a12d7c0000011159b446a5fb7ceffffffacdf3076c008d516e4d2d534541524348202a20485454502f312e310d0a486f73743a3233392e3235352e3235352e3235303a313930300d0a53543a75726e3a736368656d61732d75706e702d6f72673a6465766963653a496e7465726e6574476174657761794465766963653a310d0a4d616e3a22737364703a646973636f766572220d0a4d583a330d0a0d0a".bin
    str << "0102".bin # Tacking on a couple extra bites tht we'll strip off.
    not_stripped = UDPPacket.new
    not_stripped.read(str)
    assert_equal 135, not_stripped.udp_header.body.length
    stripped = UDPPacket.new
    stripped.read(str, :strip => true)
    assert_equal 133, stripped.udp_header.body.length
  end

  def test_udp_alter
    sample_packet = PcapFile.new.file_to_array(:f => 'sample.pcap')[0]
    pkt = Packet.parse(sample_packet)
    assert_kind_of UDPPacket, pkt
    pkt.payload = pkt.payload.gsub(/metasploit/,"MeatPistol")
    pkt.recalc
    assert_equal(0x8341, pkt.udp_sum)
    pkt.to_f('udp_test.pcap','a')
  end

  def test_udp_reread
    sample_packet = PacketFu::UDPPacket.new
    pkt = Packet.parse(sample_packet.to_s)
    assert sample_packet.is_udp?
    assert pkt.is_udp?
  end


end

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