#!/usr/bin/env ruby
# feed <emonti at matasano> 3/15/2008
#----------------------------------------------------------------------
#
# This is an eventmachine message feeder from static data sources.
# The "feed" handles messages opaquely and just plays them in the given 
# sequence.
#
# Feed can do the following things with minimum fuss: 
#   - Import messages from files, yaml, or pcap
#   - Inject custom/modified messages with "blit"
#   - Run as a server or client using UDP or TCP
#   - Bootstrap protocols without a lot of work up front
#   - Skip uninteresting messages and focus attention on the fun ones.
#   - Replay conversations for relatively unfamiliar protocols.
#   - Observe client/server behaviors using different messages at
#     various phases of a conversation.
# 
#----------------------------------------------------------------------
#
# Usage: feed -h
#
#----------------------------------------------------------------------
# To-dos / Ideas:
#  - Unix Domain Socket support?
#  - more import options?
#  - dynamic feed elements?
#  - add/control feed elements while 'stepping'?
#

require 'rubygems'
require 'eventmachine'
require 'rbkb'
require 'rbkb/plug'
require 'rbkb/plug/feed_import'
require 'rbkb/command_line'

include RBkB::CommandLine

## Default "UI" options
Plug::UI::LOGCFG[:verbose] = true
Plug::UI::LOGCFG[:dump] = :hex

## Default options sent to the Feed handler
FEED_OPTS = { 
  :close_at_end => false,
  :step => false,
  :go_first => false
}

my_addr = "0.0.0.0"
my_port = nil
listen = false
persist = false

transport = :TCP

svr_method = :start_server
cli_method = :connect


b_addr = Plug::Blit::DEFAULT_IPADDR
b_port = Plug::Blit::DEFAULT_PORT

##  Parse command-line options
arg = bkb_stdargs(nil, {})
arg.banner += " host:port"

arg.separator  "  Options:"

arg.on("-o", "--output=FILE", "Output to file") do |o|
  Plug::UI::LOGCFG[:out] = File.open(o, "w")
end

arg.on("-l", "--listen=(ADDR:?)PORT", "Server - on port (and addr?)") do |p|
  if m=/^(?:([\w\.]+):)?(\d+)$/.match(p)
    my_addr = $1 if $1
    my_port = $2.to_i
    listen = true
  else
    raise "Invalid listen argument: #{p.inspect}"
  end
end

arg.on("-b", "--blit=(ADDR:)?PORT", "Where to listen for blit") do |b|
  puts b
  unless(m=/^(?:([\w\.]+):)?(\d+)$/.match(b))
    raise "Invalid blit argument: #{b.inspect}"
  end
  b_port = m[2].to_i
  b_addr = m[1] if m[1]
end

arg.on("-i", "--[no-]initiate", "Send the first message on connect") do |i|
  FEED_OPTS[:go_first] = i
end

arg.on("-e", "--[no-]end", "End connection when feed is exhausted") do |c|
  FEED_OPTS[:close_at_end] = c
end
 
arg.on("-s", "--[no-]step", "'Continue' prompt between messages") do |s|
  FEED_OPTS[:step] = s
end

arg.on("-u", "--udp", "Use UDP instead of TCP" ) do
  transport = :UDP
  svr_method = cli_method = :open_datagram_socket
end

arg.on("-r", "--reconnect", "Attempt to reconnect endlessly.") do
  persist=true
end

arg.on("-q", "--quiet", "Suppress verbose messages/dumps") do
  Plug::UI::LOGCFG[:verbose] = false
end

arg.on("-S", "--squelch-exhausted", "Squelch 'FEED EXHAUSTED' messages") do |s|
  FEED_OPTS[:squelch_exhausted] = true
end

arg.separator  "  Sources: (can be combined)"

arg.on("-f", "--from-files=GLOB", "Import messages from raw files") do |f|
  FEED_OPTS[:feed] ||= []
  FEED_OPTS[:feed] += FeedImport.import_rawfiles(f)
end

arg.on("-x", "--from-hex=FILE", "Import messages from hexdumps") do |x|
  FEED_OPTS[:feed] ||= []
  FEED_OPTS[:feed] += FeedImport.import_dump(x)
end

arg.on("-y", "--from-yaml=FILE", "Import messages from yaml") do |y|
  FEED_OPTS[:feed] ||= []
  FEED_OPTS[:feed] += FeedImport.import_yaml(y)
end

arg.on("-p", "--from-pcap=FILE[:FILTER]", "Import messages from pcap") do |p|
  if /^([^:]+):(.+)$/.match(p)
    file = $1
    filter = $2
  else
    file = p
    filter = nil
  end

  FEED_OPTS[:feed] ||= []
  FEED_OPTS[:feed] += FeedImport.import_pcap(file, filter)
end

arg.parse!(ARGV) rescue bail "Error: #{$!}\nUse -h|--help for more info."

# parse OptionParser and import feed
begin 
  arg.parse!(ARGV) 

  # Prepare EventMachine arguments based on whether we are a client or server
  if listen
    evma_addr = my_addr
    evma_port = my_port
    meth = svr_method
    FEED_OPTS[:kind] = :server
  else

    ## Get target/listen argument for client mode
    unless (m = /^([\w\.]+):(\d+)$/.match(ARGV.shift)) and ARGV.shift.nil?
      bail arg
    end

    t_addr = m[1]
    t_port = m[2].to_i

    if transport == :UDP
      evma_addr = my_addr
      evma_port = my_port || 0
    else
      evma_addr = t_addr
      evma_port = t_port
    end

    meth = cli_method
    FEED_OPTS[:kind] = :client
  end

  FEED_OPTS[:feed] ||= []
  Plug::UI.verbose "** FEED CONTAINS #{FEED_OPTS[:feed].size} MESSAGES"

  ## error out on unexpected arguments
  raise "too many arguments" if ARGV.shift
rescue 
  bail "Error: #{$!}\nUse -h|--help for more information"
end


em_args=[ 
  meth, 
  evma_addr, 
  evma_port, 
  Plug::ArrayFeeder, 
  transport, 
  FEED_OPTS
].flatten


## Start the eventmachine
loop do
  EventMachine::run do
    EventMachine.send(*em_args) do |c|
      EventMachine.start_server(b_addr, b_port, Plug::Blit, :TCP, c)
      Plug::UI::verbose("** BLITSRV-#{b_addr}:#{b_port}(TCP) Started")

      # if this is a UDP client, we will always send the first message
      if transport == :UDP and c.kind == :client and 
        c.feed_data(c.peers.add_peer_manually(t_addr, t_port))
        c.go_first = false
      end
    end
  end

  break unless persist
  Plug::UI::verbose("** RECONNECTING")
end