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

require 'test/unit'
require 'libnet4r'

class TC_Libnet_TCP < Test::Unit::TestCase
  def test_unpack
    bytes = [
      0x43, 0x21, 0xab, 0xcd,
      0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x50, 0x0b, 0x11, 0x22,
      0x57, 0x9c, 0x00, 0x00, 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"

    t = Libnet::TCP.decode(bytes)

    assert_equal(0x4321,  t.src_port)
    assert_equal(0xabcd,  t.dst_port)
    assert_equal(32,      t.seq)
    assert_equal(0,       t.ack)
    assert_equal(5,       t.data_offset)
    assert_equal(0,       t.reserved)
    assert_equal(0x0b,    t.control)
    assert_equal(0x1122,  t.window)
    assert_equal(0x579c,  t.checksum)
    assert_equal(0,       t.urgent)
    assert_equal(payload, t.payload)

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

  def test_pack
    bytes = [
      0x43, 0x21, 0xab, 0xcd,
      0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x50, 0x0b, 0x11, 0x22,
      0x11, 0x11, 0x00, 0x00, 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"

    t = Libnet::TCP.new

    t.src_port = 0x4321
    t.dst_port = 0xabcd
    t.seq      = 32
    t.ack      = 0
    t.control  = 0x0b
    t.window   = 0x1122
    t.checksum = 0x1111
    t.urgent   = 0
    t.payload  = payload

    assert_nil(t.ptag)

    l = Libnet.new

    l.build_tcp(t)

    assert(t.ptag > 0)

    assert_equal(bytes, l.pack)

    l = Libnet.new

    t = nil
    l.build_tcp do |t|
      t.src_port = 0x4321
      t.dst_port = 0xabcd
      t.seq      = 32
      t.ack      = 0
      t.control  = 0x0b
      t.window   = 0x1122
      t.checksum = 0x1111
      t.urgent   = 0
      t.payload  = payload
    end

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

  def test_modify_ptag
    bytes1 = [
      0x43, 0x21, 0xab, 0xcd,
      0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x50, 0x0b, 0x11, 0x22,
      0x11, 0x11, 0x00, 0x00, 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, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x50, 0x0b, 0x11, 0x22,
      0x11, 0x11, 0x00, 0x00, 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

    t = l.build_tcp do |t|
      t.src_port = 0x4321
      t.dst_port = 0xabcd
      t.seq      = 32
      t.ack      = 0
      t.control  = 0x0b
      t.window   = 0x1122
      t.checksum = 0x1111
      t.urgent   = 0
      t.payload  = payload
    end

    assert(t.ptag > 0)
    ptag = t.ptag

    assert_equal(bytes1, l.pack)

    t.src_port = 0xfeed
    t.dst_port = 0xface

    l.build_tcp(t)

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

  def test_ether_ipv4_tcp
    bytes = [
      # ethernet
      0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
      0x08, 0x00,
      # ipv4
      0x45, 0x12, 0x00, 0x54, 0x43, 0x21, 0x00, 0x00, 0x48, 0x06, 0xac, 0x1b,
      0xc0, 0xa8, 0x01, 0x02, 0xc0, 0xa8, 0x01, 0x03,
      # tcp
      0x43, 0x21, 0xab, 0xcd,
      0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x50, 0x0b, 0x11, 0x22,
      0x57, 0x9c, 0x00, 0x00, 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)

    t = l.build_tcp do |t|
      t.src_port = 0x4321
      t.dst_port = 0xabcd
      t.seq      = 32
      t.ack      = 0
      t.control  = 0x0b
      t.window   = 0x1122
      t.checksum = 0
      t.urgent   = 0
      t.payload  = payload
    end

    l.build_ipv4 do |i|
      i.tos      = 0x12
      i.length   = Libnet::HL_IPV4 + Libnet::HL_TCP + payload.length
      i.id       = 0x4321
      i.frag_off = 0
      i.ttl      = 72
      i.protocol = Libnet::IPPROTO_TCP
      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
end