#!/usr/bin/env ruby

require 'icmp'
require 'getopts'

include Socket::Constants
include ICMPModule

IP_TTL = 4

class UDP < String
  def self.new(s) super(s); end
  def uh_sport; self[0..1].unpack("n")[0]; end
  def uh_dport; self[2..3].unpack("n")[0]; end
  def uh_ulen;  self[4..5].unpack("n")[0]; end
  def uh_sum;   self[6..7].unpack("n")[0]; end
end

def traceroute(host, port, nprobes, max_ttl, wait_time)
  icmp_sock = ICMPSocket.new
  udp_sock = UDPSocket.new
  ident = $$ | 0x8000
  udp_sock.bind(0, ident)
  ttl = 0
  ai = Socket::getaddrinfo(host, nil, Socket::AF_INET)[0]
  dst = ai[3]

  reached = false
  while not reached && ttl < max_ttl
    ttl += 1
    printf "%3d ", ttl
    udp_sock.setsockopt(IPPROTO_IP, IP_TTL, ttl)

    reached = false
    probed = false
    nprobes.times{|i|
      send_time = Time.now
      udp_sock.send("\0\0\0\0", 0, dst, port)
      unless ary = select([icmp_sock], nil, nil, wait_time)
        print(" *")
        next
      end
      buf = icmp_sock.recv(65535)
      recv_time = Time.now
      iph, icmpp = ICMPModule::split(buf) # IP header and ICMP packet
      ipp = icmpp.icmp_ip                 # original IP packet
      udph = UDP.new(ipp.body)            # original UDP header
      retry if ipp.ip_p != Socket::IPPROTO_UDP
      retry if ipp.ip_dst != dst || udph.uh_sport != ident

      if !probed && !reached
        ai0 = Socket::getaddrinfo(iph.ip_src, nil, Socket::AF_INET)[0]
        printf("%s (%s)", ai0[2], iph.ip_src)
      end
      printf " %.3f ms", (recv_time.to_f - send_time.to_f) * 1000

      case icmpp.icmp_type
      when ICMP_TIMXCEED
        probed = true
      when ICMP_UNREACH
        case icmpp.icmp_code
        when ICMP_UNREACH_PORT;     reached = true
        when ICMP_UNREACH_HOST;     print(" !H")
        when ICMP_UNREACH_NET;      print(" !N")
        when ICMP_UNREACH_PROTOCOL; print(" !P")
        end
      else
        print(" ??")
      end
    }
    print "\n"
  end
end

##
## Main routine
##
getopts nil, "p:33434", "q:3", "m:30", "w:3"
udp_port   = $OPT_p.to_i
nprobes    = $OPT_q.to_i
max_ttl    = $OPT_m.to_i
wait_time  = $OPT_w.to_i

$stdout.sync = true
traceroute(ARGV[0], udp_port, nprobes, max_ttl, wait_time)