# -*- coding: binary -*-
require 'packetfu/protos/tcp/reserved'
require 'packetfu/protos/tcp/hlen'
require 'packetfu/protos/tcp/ecn'
require 'packetfu/protos/tcp/flags'
require 'packetfu/protos/tcp/option'
require 'packetfu/protos/tcp/options'


module PacketFu
  # TCPHeader is a complete TCP struct, used in TCPPacket. Most IP traffic is TCP-based, by
  # volume.
  #
  # For more on TCP packets, see http://www.networksorcery.com/enp/protocol/tcp.htm
  #
  # ==== Header Definition
  # 
  #   Int16        :tcp_src       Default: random 
  #   Int16        :tcp_dst
  #   Int32        :tcp_seq       Default: random
  #   Int32        :tcp_ack
  #   TcpHlen      :tcp_hlen      Default: 5           # Must recalc as options are set. 
  #   TcpReserved  :tcp_reserved  Default: 0
  #   TcpEcn       :tcp_ecn
  #   TcpFlags     :tcp_flags
  #   Int16        :tcp_win,      Default: 0           # WinXP's default syn packet
  #   Int16        :tcp_sum,      Default: calculated  # Must set this upon generation.
  #   Int16        :tcp_urg
  #   TcpOptions   :tcp_opts
  #   String       :body
  #
  # See also TcpHlen, TcpReserved, TcpEcn, TcpFlags, TcpOpts
  class TCPHeader < Struct.new(:tcp_src, :tcp_dst,
                               :tcp_seq,
                               :tcp_ack,
                               :tcp_hlen, :tcp_reserved, :tcp_ecn, :tcp_flags, :tcp_win, 
                               :tcp_sum, :tcp_urg, 
                               :tcp_opts, :body)
    include StructFu

    def initialize(args={})
      @random_seq = rand(0xffffffff)
      @random_src = rand_port
      super(
        Int16.new(args[:tcp_src] || tcp_calc_src),
        Int16.new(args[:tcp_dst]),
        Int32.new(args[:tcp_seq] || tcp_calc_seq),
        Int32.new(args[:tcp_ack]),
        TcpHlen.new(:hlen => (args[:tcp_hlen] || 5)),
        TcpReserved.new(args[:tcp_reserved] || 0),
        TcpEcn.new(args[:tcp_ecn]),
        TcpFlags.new(args[:tcp_flags]),
        Int16.new(args[:tcp_win] || 0x4000),
        Int16.new(args[:tcp_sum] || 0),
        Int16.new(args[:tcp_urg]),
        TcpOptions.new.read(args[:tcp_opts]),
        StructFu::String.new.read(args[:body])
      )
    end

    attr_accessor :flavor

    # Helper function to create the string for Hlen, Reserved, ECN, and Flags.
    def bits_to_s
      bytes = []
      bytes[0] = (self[:tcp_hlen].to_i << 4) +
        (self[:tcp_reserved].to_i << 1) +
        self[:tcp_ecn].n.to_i
      bytes[1] = (self[:tcp_ecn].c.to_i << 7) +
        (self[:tcp_ecn].e.to_i << 6) +
        self[:tcp_flags].to_i
      bytes.pack("CC")
    end

    # Returns the object in string form.
    def to_s
      hdr = self.to_a.map do |x|
        if x.kind_of? TcpHlen
          bits_to_s
        elsif x.kind_of? TcpReserved
          next
        elsif x.kind_of? TcpEcn
          next
        elsif x.kind_of? TcpFlags
          next
        else
          x.to_s
        end
      end
      hdr.flatten.join
    end

    # Reads a string to populate the object.
    def read(str)
      force_binary(str)
      return self if str.nil?
      self[:tcp_src].read(str[0,2])
      self[:tcp_dst].read(str[2,2])
      self[:tcp_seq].read(str[4,4])
      self[:tcp_ack].read(str[8,4])
      self[:tcp_hlen].read(str[12,1])
      self[:tcp_reserved].read(str[12,1])
      self[:tcp_ecn].read(str[12,2])
      self[:tcp_flags].read(str[13,1])
      self[:tcp_win].read(str[14,2])
      self[:tcp_sum].read(str[16,2])
      self[:tcp_urg].read(str[18,2])
      self[:tcp_opts].read(str[20,((self[:tcp_hlen].to_i * 4) - 20)])
      self[:body].read(str[(self[:tcp_hlen].to_i * 4),str.size])
      self
    end

    # Setter for the TCP source port.
    def tcp_src=(i); typecast i; end
    # Getter for the TCP source port.
    def tcp_src; self[:tcp_src].to_i; end
    # Setter for the TCP destination port.
    def tcp_dst=(i); typecast i; end
    # Getter for the TCP destination port.
    def tcp_dst; self[:tcp_dst].to_i; end
    # Setter for the TCP sequence number.
    def tcp_seq=(i); typecast i; end
    # Getter for the TCP sequence number.
    def tcp_seq; self[:tcp_seq].to_i; end
    # Setter for the TCP ackowlegement number.
    def tcp_ack=(i); typecast i; end
    # Getter for the TCP ackowlegement number.
    def tcp_ack; self[:tcp_ack].to_i; end
    # Setter for the TCP window size number.
    def tcp_win=(i); typecast i; end
    # Getter for the TCP window size number.
    def tcp_win; self[:tcp_win].to_i; end
    # Setter for the TCP checksum.
    def tcp_sum=(i); typecast i; end
    # Getter for the TCP checksum.
    def tcp_sum; self[:tcp_sum].to_i; end
    # Setter for the TCP urgent field.
    def tcp_urg=(i); typecast i; end
    # Getter for the TCP urgent field.
    def tcp_urg; self[:tcp_urg].to_i; end

    # Getter for the TCP Header Length value.
    def tcp_hlen; self[:tcp_hlen].to_i; end
    # Setter for the TCP Header Length value. Can take
    # either a string or an integer. Note that if it's
    # a string, the top four bits are used.
    def tcp_hlen=(i)
      case i
      when PacketFu::TcpHlen
        self[:tcp_hlen] = i
      when Numeric
        self[:tcp_hlen] = TcpHlen.new(:hlen => i.to_i)
      else
        self[:tcp_hlen].read(i)
      end
    end

    # Getter for the TCP Reserved field.
    def tcp_reserved; self[:tcp_reserved].to_i; end
    # Setter for the TCP Reserved field.
    def tcp_reserved=(i)
      case i
      when PacketFu::TcpReserved
        self[:tcp_reserved]=i
      when Numeric
        args = {}
        args[:r1] = (i & 0b100) >> 2
        args[:r2] = (i & 0b010) >> 1
        args[:r3] = (i & 0b001)
        self[:tcp_reserved] = TcpReserved.new(args)
      else
        self[:tcp_reserved].read(i)
      end
    end

    # Getter for the ECN bits. 
    def tcp_ecn; self[:tcp_ecn].to_i; end
    # Setter for the ECN bits. 
    def tcp_ecn=(i)
      case i
      when PacketFu::TcpEcn
        self[:tcp_ecn]=i
      when Numeric
        args = {}
        args[:n] = (i & 0b100) >> 2
        args[:c] = (i & 0b010) >> 1
        args[:e] = (i & 0b001)
        self[:tcp_ecn] = TcpEcn.new(args)
      else
        self[:tcp_ecn].read(i)
      end
    end

    # Getter for TCP Options.
    def tcp_opts; self[:tcp_opts].to_s; end
    # Setter for TCP Options.
    def tcp_opts=(i)
      case i
      when PacketFu::TcpOptions
        self[:tcp_opts]=i
      else
        self[:tcp_opts].read(i)
      end
    end

    # Resets the sequence number to a new random number.
    def tcp_calc_seq; @random_seq; end
    # Resets the source port to a new random number.
    def tcp_calc_src; @random_src; end

    # Returns the actual length of the TCP options.
    def tcp_opts_len
      self[:tcp_opts].to_s.size
    end

    # Sets and returns the true length of the TCP Header.
    # TODO: Think about making all the option stuff safer. 
    def tcp_calc_hlen
      self[:tcp_hlen] = TcpHlen.new(:hlen => ((20 + tcp_opts_len) / 4))
    end

    # Generates a random high port. This is affected by packet flavor.
    def rand_port
      rand(0xffff - 1025) + 1025
    end

    # Gets a more readable option list.
    def tcp_options
     self[:tcp_opts].decode
    end

    # Gets a more readable flags list
    def tcp_flags_dotmap
      dotmap = tcp_flags.members.map do |flag|
        status = self.tcp_flags.send flag
        status == 0 ? "." : flag.to_s.upcase[0].chr
      end
      dotmap.join
    end

    # Sets a more readable option list.
    def tcp_options=(arg)
      self[:tcp_opts].encode arg
    end

    # Equivalent to tcp_src.
    def tcp_sport
      self.tcp_src.to_i
    end

    # Equivalent to tcp_src=.
    def tcp_sport=(arg)
      self.tcp_src=(arg)
    end

    # Equivalent to tcp_dst.
    def tcp_dport
      self.tcp_dst.to_i
    end
    
    # Equivalent to tcp_dst=.
    def tcp_dport=(arg)
      self.tcp_dst=(arg)
    end

    # Recalculates calculated fields for TCP (except checksum which is at the Packet level).
    def tcp_recalc(arg=:all)
      case arg
      when :tcp_hlen
        tcp_calc_hlen
      when :tcp_src
        @random_tcp_src = rand_port
      when :tcp_sport
        @random_tcp_src = rand_port
      when :tcp_seq
        @random_tcp_seq = rand(0xffffffff) 
      when :all
        tcp_calc_hlen
        @random_tcp_src = rand_port
        @random_tcp_seq = rand(0xffffffff) 
      else
        raise ArgumentError, "No such field `#{arg}'"
      end
    end

    # Readability aliases

    alias :tcp_flags_readable :tcp_flags_dotmap

    def tcp_ack_readable
      "0x%08x" % tcp_ack
    end

    def tcp_seq_readable
      "0x%08x" % tcp_seq
    end

    def tcp_sum_readable
      "0x%04x" % tcp_sum
    end

    def tcp_opts_readable
      tcp_options
    end

  end

end