require 'pcap'
require 'optparse'

def pcaplet_usage()
  $stderr.print <<END
Usage: #{File.basename $0} [ -dnv ] [ -i interface | -r file ]
       #{' ' * File.basename($0).length} [ -c count ] [ -s snaplen ] [ filter ]
Options:
  -n  do not convert address to name
  -d  debug mode
  -v  verbose mode
END
end

module Pcap
  class Pcaplet
    def usage(status, msg = nil)
      $stderr.puts msg if msg
      pcaplet_usage
      exit(status)
    end

    def initialize(args = nil)
      if args
        ARGV[0,0] = args.split(/\s+/)
      end
      @device = nil
      @rfile = nil
      @count = -1
      @snaplen = 68
      opts = OptionParser.new do |opts|
        opts.on('-d') {$DEBUG = true}
        opts.on('-v') {$VERBOSE = true}
        opts.on('-n') {Pcap.convert = false}
        opts.on('-i IFACE') {|s| @device = s}
        opts.on('-r FILE') {|s| @rfile = s}
        opts.on('-c COUNT', OptionParser::DecimalInteger) {|i| @count = i}
        opts.on('-s LEN', OptionParser::DecimalInteger) {|i| @snaplen = i}
      end
      begin
        opts.parse!
      rescue
        usage(1)
      end
      
      @filter = ARGV.join(' ')

      # check option consistency
      usage(1) if @device && @rfile
      if !@device and !@rfile
        @device = Pcap.lookupdev
      end

      # open
      begin
        if @device
          @capture = Capture.open_live(@device, @snaplen)
        elsif @rfile
          if @rfile !~ /\.gz$/
            @capture = Capture.open_offline(@rfile)
          else
            $stdin = IO.popen("gzip -dc < #@rfile", 'r')
            @capture = Capture.open_offline('-')
          end
        end
        @capture.setfilter(@filter)
      rescue PcapError, ArgumentError
        $stdout.flush
        $stderr.puts $!
        exit(1)
      end
    end

    attr('capture')

    def add_filter(f)
      if @filter == nil || @filter =~ /^\s*$/  # if empty
        @filter = f
      else
        f = f.source if f.is_a? Filter
        @filter = "( #{@filter} ) and ( #{f} )"
      end
      @capture.setfilter(@filter)
    end

    def each_packet(&block)
      begin
        duplicated = (RUBY_PLATFORM =~ /linux/ && @device == "lo")
        unless duplicated
          @capture.loop(@count, &block)
        else
          flip = true
          @capture.loop(@count) do |pkt|
            flip = (! flip)
            next if flip
            block.call pkt
          end
        end
      rescue Interrupt
        $stdout.flush
        $stderr.puts("Interrupted.")
        $stderr.puts $@.join("\n\t") if $DEBUG
      ensure
        # print statistics if live
        if @device
          stat = @capture.stats
          if stat
            $stderr.print("#{stat.recv} packets received by filter\n");
            $stderr.print("#{stat.drop} packets dropped by kernel\n");
          end
        end
      end
    end

    alias each each_packet

    def close
      @capture.close
    end
  end
end

Pcaplet = Pcap::Pcaplet