require 'active_support'

class Libnet
  class Header # :nodoc:
    include Helpers

    HWADDR_RE        = /^[A-Fa-f0-9]{2}:[A-Fa-f0-9]{2}:[A-Fa-f0-9]{2}:[A-Fa-f0-9]{2}:[A-Fa-f0-9]{2}:[A-Fa-f0-9]{2}$/
    IPADDR_RE        = /^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$/
    DEFAULT_OPTS     = {
      :units     => :bytes,
      :read_only => false,
      :optional  => false,
      :aliases   => []
    }
    UNIT_MULTIPLIERS = {
      :bits  => 1,
      :bytes => 8,
      :words => 16,
    }

    class_inheritable_accessor :assignment_filters
    class_inheritable_accessor :required_fields
    class_inheritable_accessor :optional_fields
    class_inheritable_accessor :read_only_fields

    class << self
      def inherited(cl)
        cl.assignment_filters = {}
        cl.required_fields    = []
        cl.optional_fields    = []
        cl.read_only_fields   = []
      end

      def unsigned_field(name, size, opts = {})
        opts  = DEFAULT_OPTS.merge opts
        size  = size * UNIT_MULTIPLIERS[opts[:units]]
        ivar  = "@#{name}".to_sym
        range = 0..2**size-1

        unless opts[:read_only]
          self.class_eval do
            define_method("#{name}=") do |v|
              # allow for setting to nil
              if v
                if self.assignment_filters.has_key? name
                  v = self.assignment_filters[name].call(v)
                end

                check_integer(v)

                unless range.include? v
                  raise ArgumentError, "value #{v} is out of range for field '#{name}'"
                end
              end

              instance_variable_set(ivar, v)
            end
          end
        end

        add_field(name, opts)

        nil
      end

      def octets_field(name, size = nil, opts = {})
        opts = DEFAULT_OPTS.merge opts
        ivar = "@#{name}".to_sym

        unless opts[:read_only]
          self.class_eval do
            define_method("#{name}=") do |v|
              if v
                value = v.dup
                if self.assignment_filters.has_key? name
                  value = self.assignment_filters[name].call(v)
                end

                check_string(value)

                if size && value.length != size
                  raise ArgumentError, "value must be #{size} bytes long for field '#{name}'"
                end
              else
                # allow for setting to nil
                value = v
              end

              instance_variable_set(ivar, value)
            end
          end
        end

        add_field(name, opts)

        nil
      end

      def assignment_filter(*fields, &blk)
        fields.each do |f|
          self.assignment_filters[f] = blk
        end
      end

      def decode(bytes)
        raise NotImplementedError, "class #{self} has not implemented a decode method"
      end

      def fields
        self.required_fields + self.optional_fields + self.read_only_fields
      end

      def add_field(name, opts)
        ivar = "@#{name}".to_sym

        # add this field to the list
        if opts[:read_only]
          self.read_only_fields << name
        elsif opts[:optional]
          self.optional_fields << name
        else
          self.required_fields << name
        end

        self.class_eval do
          # create the query method
          define_method("#{name}?") do
            !!instance_variable_get(ivar)
          end

          # create the reader method
          define_method(name) do
            instance_variable_get(ivar)
          end

          # create any aliases
          opts[:aliases].each do |a|
            alias_method a, name
            alias_method "#{a}?".to_sym, "#{name}?".to_sym

            unless opts[:read_only]
              alias_method "#{a}=".to_sym, "#{name}=".to_sym
            end
          end

          # add an auto-checksum method if this is a checksum field
          if name == :checksum
            define_method(:auto_checksum?) do |bool|
              instance_variable_set(:@auto_checksum)
            end

            define_method(:auto_checksum=) do |bool|
              instance_variable_set(:@auto_checksum, bool)
            end
          end
        end
      end
    end

    private_class_method :add_field
    attr_reader          :ptag

    def initialize(bytes = nil)
      self.class.decode(bytes) if bytes

      # only used if class has a checksum field
      @auto_checksum = true

      self
    end

    private

    def check_packable
      not_set = []

      self.required_fields.each do |f|
        not_set << f unless self.send("#{f}?")
      end

      if not_set.length > 0
        raise ArgumentError, "the following #{self.class} required fields have not been set: " + not_set.join(", ")
      end

      true
    end
  end

  # = Ethernet
  #
  # Class to represent an IEEE 802.3 Ethernet header.
  #
  # === Example
  #
  #   e = Libnet::Ethernet.new
  #
  #   e.dst     = 'aa:bb:cc:dd:ee:ff'
  #   e.src     = [0x11, 0x22, 0x33, 0x44, 0x55, 0x66].pack("C*")
  #   e.type    = 0x1234
  #   e.payload = "this is the message"
  #
  # == Public Class Methods
  #
  # === decode(string) -> new Ethernet object
  #
  # Decode a packed ethernet packet.
  #
  # == Public Instance Methods
  #
  # === dst, dst=, dst?
  # Get/set/query the destination MAC address.
  #
  # === src, src=, src?
  # Get/set/query the source MAC address.
  #
  # === type, type=, type?
  # Get/set/query the protocol type.
  #
  # === payload, payload=, payload?
  # Get/set/query the ethernet payload. Optional.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class Ethernet < Header
    octets_field   :dst,     6
    octets_field   :src,     6
    unsigned_field :type,    2
    octets_field   :payload, nil, :optional => true

    assignment_filter(:dst, :src) do |v|
      if v.match(HWADDR_RE)
        Libnet.hex_aton(v)
      else
        v
      end
    end
  end

  # = VLAN
  #
  # Class to represent an IEEE 802.1q VLAN tagging header.
  #
  # === Example
  #
  #   v = Libnet::VLAN.new
  #
  #   v.dst      = 'aa:bb:cc:dd:ee:ff'
  #   v.src      = '11:22:33:44:55:66'
  #   v.tpi      = 0xbeef
  #   v.priority = 6
  #   v.cfi      = 0
  #   v.id       = 2730
  #   v.type     = 0x0800
  #   v.payload  = "this is the message"
  #
  # == Public Class Methods
  #
  # === decode(string) -> new VLAN object
  #
  # Decode a packed vlan packet.
  #
  # == Public Instance Methods
  #
  # === dst, dst=, dst?
  # Get/set/query the destination MAC address.
  #
  # === src, src=, src?
  # Get/set/query the source MAC address.
  #
  # === tpi, tpi=, tpi?
  # Get/set/query the tag protocol identifier.
  #
  # === priority, priority=, priority?
  # Get/set/query the priority.
  #
  # === cfi, cfi=, cfi?
  # Get/set/query the canonical format indicator.
  #
  # === id, id=, id?
  # Get/set/query the vlan identifier.
  #
  # === len, len=, len? or type, type=, type?
  # Get/set/query the length/type.
  #
  # === payload, payload=, payload?
  # Get/set/query the payload.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class VLAN < Header
    octets_field   :dst,       6
    octets_field   :src,       6
    unsigned_field :tpi,       2
    unsigned_field :priority,  3, :units => :bits
    unsigned_field :cfi,       1, :units => :bits
    unsigned_field :id,       12, :units => :bits
    unsigned_field :len,       2, :aliases => [ :type ]
    octets_field   :payload,  nil, :optional => true

    assignment_filter(:dst, :src) do |v|
      if v.match(HWADDR_RE)
        Libnet.hex_aton(v)
      else
        v
      end
    end
  end

  # = IPv4
  #
  # Class to represent a version 4 RFC 791 Internet Protocol (IP) header.
  #
  # === Example
  #
  #   i = Libnet::IPv4.new
  #
  #   i.tos      = 0x12
  #   i.length   = Libnet::HL_IPV4 + 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
  #   i.payload  = "this is the message"
  #
  # == Public Class Methods
  #
  # === decode(string) -> new IPv4 object
  #
  # Decode a packed IPv4 packet.
  #
  # == Public Instance Methods
  #
  # === version, version?
  # Get/query the IP version.
  #
  # === ihl, ihl?
  # Get/query the IP header length.
  #
  # === tos, tos=, tos?
  # Get/set/query the type of service.
  #
  # === length, length=, length?
  # Get/set/query the total length of the IP packet.
  #
  # === id, id=, id?
  # Get/set/query the IP identification number.
  #
  # === frag_off, frag_off=, frag_off?
  # Get/set/query the fragmentation bits and offset.
  #
  # === ttl, ttl=, ttl?
  # Get/set/query the time to live.
  #
  # === protocol, protocol=, protocol?
  # Get/set/query the upper layer protocol.
  #
  # === checksum, checksum=, checksum?
  # Get/set/query the checksum.  Set to 0 for Libnet to autofill.
  #
  # === src_ip, src_ip=, src_ip?
  # Get/set/query the source IP address.
  #
  # === dst_ip, dst_ip=, dst_ip?
  # Get/set/query the destination IP address.
  #
  # === payload, payload=, payload?
  # Get/set/query the payload.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class IPv4 < Header
    unsigned_field :version,   4, :units => :bits, :read_only => true
    unsigned_field :ihl,       4, :units => :bits, :read_only => true
    unsigned_field :tos,       1
    unsigned_field :length,    2
    unsigned_field :id,        2
    unsigned_field :frag_off,  2
    unsigned_field :ttl,       1
    unsigned_field :protocol,  1
    unsigned_field :checksum,  2
    unsigned_field :src_ip,    4
    unsigned_field :dst_ip,    4
    octets_field   :payload, nil, :optional => true

    assignment_filter(:src_ip, :dst_ip) do |v|
      if v.is_a? String
        Libnet.ipv4_aton(v)
      else
        v
      end
    end
  end

  # = IPv6
  #
  # Class to represent a version 6 RFC 2460 Internet Protocol (IP) header.
  #
  # === Example
  #
  #   i = Libnet::IPv6.new
  #
  #   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"
  #   i.payload       = "this is the message"
  #
  # == Public Class Methods
  #
  # === decode(string) -> new IPv6 object
  #
  # Decode a packed IPv6 packet.
  #
  # == Public Instance Methods
  #
  # === version, version?
  # Get/query the IP version.
  #
  # === traffic_class, traffic_class=, traffic_class?
  # Get/set/query the traffic class.
  #
  # === flow_label, flow_label=, flow_label?
  # Get/set/query the flow label.
  #
  # === length, length=, length?
  # Get/set/query the length of the IPv6 packet minus the 40 byte IPv6 header.
  #
  # === next_header, next_header=, next_header?
  # Get/set/query the next header.
  #
  # === hop_limit, hop_limit=, hop_limit?
  # Get/set/query the hop limit.
  #
  # === src_ip, src_ip=, src_ip?
  # Get/set/query the source IPv6 address.
  #
  # === dst_ip, dst_ip=, dst_ip?
  # Get/set/query the destination IPv6 address.
  #
  # === payload, payload=, payload?
  # Get/set/query the payload.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class IPv6 < Header
    unsigned_field :version,        4, :units => :bits, :read_only => true
    unsigned_field :traffic_class,  1
    unsigned_field :flow_label,    20, :units => :bits
    unsigned_field :length,         2
    unsigned_field :next_header,    1
    unsigned_field :hop_limit,      1
    octets_field   :src_ip,        16
    octets_field   :dst_ip,        16
    octets_field   :payload, nil, :optional => true

    assignment_filter(:src_ip, :dst_ip) do |v|
      if v.length != 16
        Libnet.ipv6_aton(v)
      else
        v
      end
    end
  end

  # = UDP
  #
  # Class to represent an RFC 768 User Datagram Protocol (UDP) header.
  #
  # === Example
  #
  #   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
  #
  # == Public Class Methods
  #
  # === decode(string) -> new UDP object
  #
  # Decode a packed UDP packet.
  #
  # == Public Instance Methods
  #
  # === src_port, src_port=, src_port?
  # Get/set/query the source port.
  #
  # === dst_port, dst_port=, dst_port?
  # Get/set/query the destination port.
  #
  # === length, length=, length?
  # Get/set/query the total length of the UDP packet.
  #
  # === checksum, checksum=, checksum?
  # Get/set/query the checksum.  Set to 0 for Libnet to autofill.
  #
  # === payload, payload=, payload?
  # Get/set/query the payload.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class UDP < Header
    unsigned_field :src_port, 2
    unsigned_field :dst_port, 2
    unsigned_field :length,   2
    unsigned_field :checksum, 2
    octets_field   :payload, nil, :optional => true
  end

  # = TCP
  #
  # Class to represent an RFC 793 Transmission Control Protocol (TCP) header.
  #
  # === Example
  #
  #   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
  #
  # == Public Class Methods
  #
  # === decode(string) -> new TCP object
  #
  # Decode a packed TCP packet.
  #
  # == Public Instance Methods
  #
  # === src_port, src_port=, src_port?
  # Get/set/query the source port.
  #
  # === dst_port, dst_port=, dst_port?
  # Get/set/query the destination port.
  #
  # === seq, seq=, seq?
  # Get/set/query the sequence number.
  #
  # === ack, ack=, ack?
  # Get/set/query the acknowledgement number.
  #
  # === data_offset, data_offset?
  # Get/query the data offset.
  #
  # === reserved, reserved?
  # Get/query the reserved field.
  #
  # === control, control=, control?
  # Get/set/query the control flags.
  #
  # === window, window=, window?
  # Get/set/query the window size.
  #
  # === checksum, checksum=, checksum?
  # Get/set/query the checksum.
  #
  # === urgent, urgent=, urgent?
  # Get/set/query the urgent pointer.
  #
  # === payload, payload=, payload?
  # Get/set/query the payload.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class TCP < Header
    unsigned_field :src_port,    2
    unsigned_field :dst_port,    2
    unsigned_field :seq,         4
    unsigned_field :ack,         4
    unsigned_field :data_offset, 4, :units => :bits, :read_only => true
    unsigned_field :reserved,    6, :units => :bits, :read_only => true
    unsigned_field :control,     6, :units => :bits
    unsigned_field :window,      2
    unsigned_field :checksum,    2
    unsigned_field :urgent,      2
    octets_field   :payload, nil, :optional => true
  end

  # = ARP
  #
  # Class to represent an Address Resolution Protocol (ARP) header.  Depending
  # on the +operation+ value, the method builds one of several different types
  # of RFC 826 or RFC 903 RARP packets.
  #
  # === Example
  #
  #   a = Libnet::ARP.new
  #
  #   a.htype        = 1
  #   a.protocol     = 0x0800
  #   a.hlen         = 6
  #   a.plen         = 4
  #   a.operation    = 1
  #   a.sender_haddr = 'aa:bb:cc:dd:ee:ff'
  #   a.sender_paddr = '192.168.1.100'
  #   a.target_haddr = '11:22:33:44:55:66'
  #   a.target_paddr = '192.168.1.123'
  #
  # == Public Class Methods
  #
  # === decode(string) -> new ARP object
  #
  # Decode a packed ARP packet.
  #
  # == Public Instance Methods
  #
  # === htype, htype=, htype?
  # Get/set/query the hardware address type.
  #
  # === protocol, protocol=, protocol?
  # Get/set/query the protocol address type.
  #
  # === hlen, hlen=, hlen?
  # Get/set/query the hardware address length.
  #
  # === plen, plen=, plen?
  # Get/set/query the protocol address length.
  #
  # === operation, operation=, operation?
  # Get/set/query the ARP operation type.
  #
  # === sender_haddr, sender_haddr=, sender_haddr?
  # Get/set/query the sender's hardware address.
  #
  # === sender_paddr, sender_paddr=, sender_paddr?
  # Get/set/query the sender's protocol address.
  #
  # === target_haddr, target_haddr=, target_haddr?
  # Get/set/query the target's hardware address.
  #
  # === target_paddr, target_paddr=, target_paddr?
  # Get/set/query the target' protocol address.
  #
  # === payload, payload=, payload?
  # Get/set/query the payload.
  #
  # === ptag
  # Get the ptag value for this header object.  Set by the Libnet builder method.
  #
  class ARP < Header
    unsigned_field :htype,     2
    unsigned_field :protocol,  2
    unsigned_field :hlen,      1
    unsigned_field :plen,      1
    unsigned_field :operation, 2
    octets_field   :sender_haddr
    octets_field   :sender_paddr
    octets_field   :target_haddr
    octets_field   :target_paddr
    octets_field   :payload,      nil, :optional => true

    assignment_filter(:sender_haddr, :target_haddr) do |v|
      if v.match(HWADDR_RE)
        Libnet.hex_aton(v)
      else
        v
      end
    end

    assignment_filter(:sender_paddr, :target_paddr) do |v|
      if v.match(IPADDR_RE)
        ip = Libnet.ipv4_aton(v)
        [ ip ].pack("N")
      else
        v
      end
    end
  end
end