$:.unshift File.join(File.dirname(__FILE__), "..", "lib")

require 'test/unit'
require 'libnet4r'

class TC_Libnet_UDP < Test::Unit::TestCase
  def test_unpack
    bytes = [
      0x43, 0x21, 0xab, 0xcd, 0x00, 0x34, 0x11, 0x11,
      0x74, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63, 0x6b, 0x20, 0x62, 0x72,
      0x6f, 0x77, 0x6e, 0x20, 0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
      0x65, 0x64, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20,
      0x6c, 0x61, 0x7a, 0x79, 0x20, 0x64, 0x6f, 0x67
    ].pack("C*")
    payload = "the quick brown fox jumped over the lazy dog"

    u = Libnet::UDP.decode(bytes)

    assert_equal(0x4321, u.src_port)
    assert_equal(0xabcd, u.dst_port)
    assert_equal(0x0034, u.length)
    assert_equal(0x1111, u.checksum)
    assert_equal(payload, u.payload)

    assert_raise ArgumentError do
      Libnet::UDP.decode("foobar")
    end
  end

  def test_pack
    bytes = [
      0x43, 0x21, 0xab, 0xcd, 0x00, 0x34, 0x11, 0x11,
      0x74, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63, 0x6b, 0x20, 0x62, 0x72,
      0x6f, 0x77, 0x6e, 0x20, 0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
      0x65, 0x64, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20,
      0x6c, 0x61, 0x7a, 0x79, 0x20, 0x64, 0x6f, 0x67
    ].pack("C*")
    payload = "the quick brown fox jumped over the lazy dog"

    u = Libnet::UDP.new

    u.src_port = 0x4321
    u.dst_port = 0xabcd
    u.length   = Libnet::HL_UDP + payload.length
    u.checksum = 0x1111
    u.payload  = payload

    assert_nil(u.ptag)

    l = Libnet.new

    l.build_udp(u)

    assert(u.ptag > 0)
    assert_equal(bytes, l.pack)

    l = Libnet.new

    u = nil
    l.build_udp do |u|
      u.src_port = 0x4321
      u.dst_port = 0xabcd
      u.length   = Libnet::HL_UDP + payload.length
      u.checksum = 0x1111
      u.payload  = payload
    end

    assert(u.ptag > 0)
    assert_equal(bytes, l.pack)
  end

  def test_modify_ptag
    bytes1 = [
      0x43, 0x21, 0xab, 0xcd, 0x00, 0x34, 0x11, 0x11,
      0x74, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63, 0x6b, 0x20, 0x62, 0x72,
      0x6f, 0x77, 0x6e, 0x20, 0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
      0x65, 0x64, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20,
      0x6c, 0x61, 0x7a, 0x79, 0x20, 0x64, 0x6f, 0x67
    ].pack("C*")
    bytes2 = [
      0xfe, 0xed, 0xfa, 0xce, 0x00, 0x34, 0x11, 0x11,
      0x74, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63, 0x6b, 0x20, 0x62, 0x72,
      0x6f, 0x77, 0x6e, 0x20, 0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
      0x65, 0x64, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20,
      0x6c, 0x61, 0x7a, 0x79, 0x20, 0x64, 0x6f, 0x67
    ].pack("C*")
    payload = "the quick brown fox jumped over the lazy dog"

    l = Libnet.new

    u = l.build_udp do |u|
      u.src_port = 0x4321
      u.dst_port = 0xabcd
      u.length   = Libnet::HL_UDP + payload.length
      u.checksum = 0x1111
      u.payload  = payload
    end

    assert(u.ptag > 0)
    ptag = u.ptag
    assert_equal(bytes1, l.pack)

    u.src_port = 0xfeed
    u.dst_port = 0xface

    l.build_udp(u)

    assert_equal(ptag, u.ptag)
    assert_equal(bytes2, l.pack)
  end

  def test_ether_ipv4_udp
    bytes = [
      # ethernet
      0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
      0x08, 0x00,
      # ipv4
      0x45, 0x12, 0x00, 0x48, 0x43, 0x21, 0x00, 0x00, 0x48, 0x11, 0xac, 0x1c,
      0xc0, 0xa8, 0x01, 0x02, 0xc0, 0xa8, 0x01, 0x03,
      # udp
      0x43, 0x21, 0xab, 0xcd,
      0x00, 0x34, 0xb8, 0xb6, 0x74, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63,
      0x6b, 0x20, 0x62, 0x72, 0x6f, 0x77, 0x6e, 0x20, 0x66, 0x6f, 0x78, 0x20,
      0x6a, 0x75, 0x6d, 0x70, 0x65, 0x64, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20,
      0x74, 0x68, 0x65, 0x20, 0x6c, 0x61, 0x7a, 0x79, 0x20, 0x64, 0x6f, 0x67
    ].pack("C*")
    payload = "the quick brown fox jumped over the lazy dog"

    l = Libnet.new(:link)

    u = l.build_udp do |u|
      u.src_port = 0x4321
      u.dst_port = 0xabcd
      u.length   = Libnet::HL_UDP + payload.length
      u.checksum = 0
      u.payload  = payload
    end

    l.build_ipv4 do |i|
      i.tos      = 0x12
      i.length   = Libnet::HL_IPV4 + Libnet::HL_UDP + payload.length
      i.id       = 0x4321
      i.frag_off = 0
      i.ttl      = 72
      i.protocol = Libnet::IPPROTO_UDP
      i.checksum = 0
      i.src_ip   = '192.168.1.2'
      i.dst_ip   = 0xc0a80103
    end

    l.build_ethernet do |e|
      e.dst = 'aa:bb:cc:dd:ee:ff'
      e.src = '11:22:33:44:55:66'
      e.type = 0x0800
    end

    assert_equal(bytes, l.pack)

    if Process.uid > 0 && Process.euid > 0
      # user is not root, so we should get an exception here
      assert_raise RuntimeError do
        l.write
      end
    else
      # user is root, so this should work
      assert_equal(bytes.length, l.write)
    end
  end

  def test_ether_ipv6_udp
    bytes = [
      # ethernet
      0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
      0x86, 0xdd,
      # ipv6
      0x62, 0x2d, 0xbe, 0xef, 0x00, 0x54, 0x11, 0x18, 0x00, 0x11, 0x22, 0x33,
      0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,
      0xaa, 0xbb, 0xaa, 0xbb, 0xcc, 0xdd, 0xcc, 0xdd, 0xee, 0xff, 0xee, 0xff,
      0x22, 0x33, 0x22, 0x33,
      # udp
      0xfe, 0xed, 0xfa, 0xce, 0x00, 0x34, 0x63, 0x62,
      0x74, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63, 0x6b, 0x20, 0x62, 0x72,
      0x6f, 0x77, 0x6e, 0x20, 0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
      0x65, 0x64, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20,
      0x6c, 0x61, 0x7a, 0x79, 0x20, 0x64, 0x6f, 0x67  
    ].pack("C*")
    payload = "the quick brown fox jumped over the lazy dog"

    l = Libnet.new(:link)

    u = l.build_udp do |u|
      u.src_port = 0xfeed
      u.dst_port = 0xface
      u.length   = Libnet::HL_UDP + payload.length
      u.checksum = 0
      u.payload  = payload
    end

    l.build_ipv6 do |i|
      i.traffic_class = 0x22
      i.flow_label    = 0xdbeef
      i.length        = Libnet::HL_IPV6 + payload.length
      i.next_header   = Libnet::IPPROTO_UDP
      i.hop_limit     = 24
      i.src_ip        = [
        0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb,
        0xcc, 0xdd, 0xee, 0xff
      ].pack("C*")
      i.dst_ip        = "aabb:aabb:ccdd:ccdd:eeff:eeff:2233:2233"
    end

    l.build_ethernet do |e|
      e.dst = 'aa:bb:cc:dd:ee:ff'
      e.src = '11:22:33:44:55:66'
      e.type = 0x86dd
    end

    assert_equal(bytes, l.pack)
    if Process.uid > 0 && Process.euid > 0
      # user is not root, so we should get an exception here
      assert_raise RuntimeError do
        l.write
      end
    else
      # user is root, so this should work
      assert_equal(bytes.length, l.write)
    end
  end
end