require 'socket'
require 'ipaddr'

class RawSocket # :nodoc:
  @@id_arr = []

  def initialize(src_addr, dest_addr)
    # Define socket
    begin
      @socket = Socket.new PF_INET, SOCK_RAW, IPPROTO_RAW
    rescue SystemCallError => e
      raise SystemCallError, "You must be root to use raw sockets! #{e}"
    end

    @socket.setsockopt IPPROTO_IP, IP_HDRINCL, 1

    # Checks addresses
    @src_addr  = check_addr src_addr
    @dest_addr = check_addr dest_addr

    # Source and destination port are zero
    @src_port  = 0
    @dest_port = 0

    # Set correct protocol version in the header
    @version = @dest_addr.ipv4? ? "0100" : "0110"

    # Total lenght: must be overridden by subclasses
    @tot_lenght = 20

    # Protocol: must be overridden by subclasses
    @protocol = 1 # ICMP by default

    # Generate a new id
    # @id = genID
    @id = 1234

    # Generate peer sockaddr
    @to = Socket.pack_sockaddr_in @dest_port, @dest_addr.to_s
  end

  def send(payload = '')
    packet = make_ip_header([
      [@version + '0101', 'B8'],          # version, hlen
      [0, 'C'],                           # tos
      [@tot_lenght + payload.size, 'n'],  # total len
      [@id, 'n'],                         # id
      [0, 'n'],                           # flags, offset
      [64, 'C'],                          # ttl
      [@protocol, 'C'],                   # protocol
      [0, 'n'],                           # checksum
      [@src_addr.to_i, 'N'],              # source
      [@dest_addr.to_i, 'N'],             # destination
    ])
    packet << make_transport_header(payload.size)
    packet << [payload].pack("a*")
    @socket.send(packet, 0, @to)
  end

  private

  def check_addr(addr)
    case addr
    when String
      IPAddr.new(addr)
    when IPAddr
      addr
    else
      raise ArgumentError, "Wrong address format: #{addr}"
    end
  end

  def check_port(port)
    if (1..65_535).cover?(port) && port.is_a?(Integer)
      port
    else
      raise ArgumentError, "Port #{port} not valid"
    end
  end

  def genID
    while @@id_arr.include?(q = rand(65_535))
    end
    @@id_arr.push(q)
    q
  end

  def ipchecksum(data)
    checksum = data.unpack("n*").inject(0) { |s, x| s + x }
    ((checksum >> 16) + (checksum & 0xffff)) ^ 0xffff
  end

  def make_ip_header(parts)
    template = ''
    data = []
    parts.each do |part|
      data += part[0..-2]
      template << part[-1]
    end
    data_str = data.pack(template)
    checksum = ipchecksum(data_str)
    data[-3] = checksum
    data.pack(template)
  end

  def make_transport_header
    ""
  end
end

class UdpRawSocket < RawSocket # :nodoc:
  def initialize(src_addr, src_port, dest_addr, dest_port)
    super(src_addr, dest_addr)

    # Check ports
    @src_port  = check_port src_port
    @dest_port = check_port dest_port

    # Total lenght: must be overridden by subclasses
    @tot_lenght = 20 + 8 # 8 bytes => UDP Header

    # Protocol: must be overridden by subclasses
    @protocol = 17 # UDP protocol

    @to = Socket.pack_sockaddr_in @dest_port, @dest_addr.to_s
  end

  private

  def make_udp_header(parts)
    template = ''
    data = []
    parts.each do |part|
      data += part[0..-2]
      template << part[-1]
    end
    data.pack(template)
  end

  def make_transport_header(pay_size)
    make_udp_header([
      [@src_port, 'n'], # source port
      [@dest_port, 'n'],       # destination port
      [8 + pay_size, 'n'],     # len
      [0, 'n']                 # checksum (mandatory)
    ])
  end
end