require 'rbkb/cli'
require 'rbkb/plug'
require 'rbkb/plug/feed_import'
require 'eventmachine'

# Copyright 2009 emonti at moatasano dot com
# See README.rdoc for license information
#
# This is a plug-board 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.
#
# To-dos / Ideas:
#  - Unix Domain Socket support?
#  - more import options?
#  - dynamic feed elements?
#  - add/control feed elements while 'stepping'?
#

class Rbkb::Cli::Feed < Rbkb::Cli::Executable
  def initialize(*args)
    @local_addr = '0.0.0.0'
    @local_port = nil
    @listen = false
    @persist = false
    @transport = :TCP
    @svr_method = :start_server
    @cli_method = :bind_connect
    @blit_addr = Plug::Blit::DEFAULT_IPADDR
    @blit_port = Plug::Blit::DEFAULT_PORT

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

    super(*args)

    # TODO: Plug::UI obviously need fixing.
    # TODO It shouldn't be driven by constants for configuration
    Plug::UI::LOGCFG[:verbose] = true
    Plug::UI::LOGCFG[:dump] = :hex
    Plug::UI::LOGCFG[:out] = @stderr
  end

  def make_parser
    arg = super()
    arg.banner += ' host:port'

    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|
      raise "Invalid listen argument: #{p.inspect}" unless /^(?:([\w._-]+):)?(\d+)$/.match(p)

      @local_addr = ::Regexp.last_match(1) if ::Regexp.last_match(1)
      @local_port = ::Regexp.last_match(2).to_i
      @listen = true
    end

    arg.on('-s', '--source=(ADDR:?)PORT', 'Bind client on port and addr') do |p|
      if /^(?:([\w.]+):)?(\d+)$/.match(p)
        @local_addr = ::Regexp.last_match(1) if ::Regexp.last_match(1)
        @local_port = ::Regexp.last_match(2).to_i
      else
        bail("Invalid source 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

      @blit_port = m[2].to_i
      @blit_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('--[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
    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('-Q', '--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 = ::Regexp.last_match(1)
        filter = ::Regexp.last_match(2)
      else
        file = p
        filter = nil
      end

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

  def parse(*args)
    super(*args)

    @svr_method = @cli_method = :open_datagram_socket if @transport == :UDP

    @local_port ||= 0
    # Prepare EventMachine arguments based on whether we are a client or server
    if @listen # server
      @meth = @svr_method
      addr_args = [@local_addr, @local_port]
      @feed_opts[:kind] = :server
      @feed_opts[:no_stop_on_unbind] = true
    else # client

      ## Get target/listen argument for client mode
      unless (m = /^([\w.]+):(\d+)$/.match(tgt = @argv.shift))
        bail_args tgt
      end

      @target_addr = m[1]
      @target_port = m[2].to_i

      addr_args = if @transport == :UDP
                    [@local_addr, @local_port]
                  else
                    [@local_addr, @local_port, @target_addr, @target_port]
                  end

      @meth = @cli_method
      @feed_opts[:kind] = :client
    end

    @feed_opts[:feed] ||= []

    @em_args = [
      @meth,
      addr_args,
      Plug::ArrayFeeder,
      @transport,
      @feed_opts
    ].flatten

    parse_catchall
  end

  def go(*args)
    super(*args)

    Plug::UI.verbose "** FEED CONTAINS #{@feed_opts[:feed].size} MESSAGES"

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

          # if this is a UDP client, we will always send the first message
          if %i[UDP client] == [@transport, c.kind]
            peer = c.peers.add_peer_manually(@target_addr, @target_port)
            c.feed_data(peer)
            c.go_first = false
          end
        end
      end

      break unless @persist

      Plug::UI.verbose('** RECONNECTING')
    end
  end
end
