#!/usr/bin/env ruby
# Copyright 2009 emonti at matasano.com
# See README.rdoc for license information
#
# A blit-able reverse TCP proxy. Displays traffic hexdumps. Currently uses
# the default blit port for its blit receiver.
#
# XXX TODO - refactor me!

begin
  require 'rubygems'
rescue LoadError
end
require 'eventmachine'
require 'socket'
require 'optparse'
require 'rbkb/plug'

class BlitPlug
  module UI
    def log(*msg)
      return if PLUG_OPTS[:quiet]

      PLUG_OPTS[:out].puts msg
    end
    module_function :log
  end

  class Controller
    attr_accessor :tgtaddr, :tgtport, :tgtclient, :blit, :peers

    @@controller = nil

    def initialize(tgtaddr, tgtport, tgtclient)
      @tgtaddr = tgtaddr
      @tgtport = tgtport
      @tgtclient = tgtclient

      @@controller = self

      @peers = []

      ## Just tack on a blit server???
      @blit = EventMachine.start_server(
        PLUG_OPTS[:blit_addr], PLUG_OPTS[:blit_port], Plug::Blit, :TCP, self
      )
    end

    # #----------------------------------------

    def dispatch_rcv(_snder, data)
      data # for now
    end

    # #----------------------------------------

    def dispatch_close(_snder)
      nil # for now
    end

    # #----------------------------------------

    def self.proxy(cli)
      unless (ctrl = @@controller)
        raise "No controller exists for this connection: #{cli.sock_peername}"
      end

      tgtaddr = ctrl.tgtaddr
      tgtport = ctrl.tgtport
      tgtclient = ctrl.tgtclient

      srv = EventMachine.connect(tgtaddr, tgtport, tgtclient)
      srv.plug_peers.push cli
      cli.plug_peers.push srv

      ctrl.peers.push srv
      ctrl.peers.push cli ### I suppose this is probably useful too..

      srv.controller = cli.controller = ctrl
    end
  end # class BlitPlug::Controller

  module BaseTCP
    include UI

    attr_accessor :plug_peers, :controller, :kind
    attr_reader :sock_peer, :sock_peername

    def post_init
      @plug_peers = []
      @kind = :conn # default
    end

    def name
      @name
    end

    def say(data, sender)
      log "%#{sender.kind.to_s.upcase}-SAYS", data.hexdump(out: StringIO.new), '%'
      send_data data
    end

    def receive_data(data)
      log "%#{kind.to_s.upcase}-#{sock_peername}-SAYS", data.hexdump, '%'
      if @controller and (data = @controller.dispatch_rcv(self, data)).nil?
        return
      end

      @plug_peers.each { |p| p.send_data data }
    end

    def notify_connection
      @name = "#{kind.to_s.upcase}-#{sock_peername}"
      log "%#{@name}-CONNECTED"
    end

    def unbind
      @name = "#{kind.to_s.upcase}-#{sock_peername}"
      log "%#{@name}-CLOSED"

      cret = (@controller and @controller.dispatch_close(self))

      @plug_peers.each do |p|
        p.plug_peers.delete(self)
        p.close_connection unless cret
      end
    end
  end

  module TCPListener
    include BlitPlug::BaseTCP
    attr_accessor :tgtaddr, :tgtport

    def post_init
      super
      @kind = :client
      @sock_peer = Socket.unpack_sockaddr_in(get_peername).reverse
      @sock_peername = @sock_peer.join(':')

      @controller = BlitPlug::Controller.proxy(self)

      start_tls if PLUG_OPTS[:svr_tls]

      notify_connection
    end
  end # module TCPListener

  module TCPClient
    include BlitPlug::BaseTCP
    attr_accessor :connected

    def post_init
      super
      @kind = :server
    end

    def connection_completed
      @sock_peer = Socket.unpack_sockaddr_in(get_peername).reverse
      @sock_peername = @sock_peer.join(':')
      notify_connection
      start_tls if PLUG_OPTS[:tgt_tls]
    end
  end # module TCPClient
end # module BlitPlug

PLUG_OPTS = {
  quiet: false,
  out: STDOUT,
  blit_addr: Plug::Blit::DEFAULT_IPADDR,
  blit_port: Plug::Blit::DEFAULT_PORT
}

def bail(*msg)
  warn msg
  exit 1
end

#############################################################################
### MAIN
#############################################################################
#
# Get option arguments
opts = OptionParser.new do |opts|
  opts.banner = "Usage: #{File.basename $0} [options] target:tport[@[laddr:]lport]\n",
                "  <target:tport>  = the address of the target service\n",
                "  <@laddr:lport> = optional address and port to listen on\n"

  opts.separator ''
  opts.separator 'Options:'

  opts.on_tail('-h', '--help', 'Show this message') do
    puts opts
    exit 1
  end

  opts.on('-o', '--output FILE', 'send output to a file') do |o|
    PLUG_OPTS[:out] = begin
      File.open(o, 'w')
    rescue StandardError
      (bail $!)
    end
  end

  opts.on('-l', '--listen ADDR:PORT',
          'optional listener address:port',
          '(default: 0.0.0.0:<tport>)') do |addr|
    unless m = /^([\w.]+)?(?::(\d+))?$/.match(addr)
      warn 'invalid listener address'
      exit 1
    end
    PLUG_OPTS[:svraddr] = m[1]
    PLUG_OPTS[:svrport] = m[2] ? m[2].to_i : nil
  end

  opts.on('-q', '--[no-]quiet', 'Suppress/Enable conversation dumps.') do |q|
    PLUG_OPTS[:quiet] = q
  end

  opts.on('-b', '--blitsrv ADDR:PORT',
          'specify blit listener [address:]port',
          "(default: #{PLUG_OPTS[:blit_addr]}:#{PLUG_OPTS[:blit_port]})") do |addr|
    unless m = /^(?:([\w.]+):)?(\d+)$/.match(addr)
      warn 'invalid blit listener argument'
      exit 1
    end
    PLUG_OPTS[:blit_addr] = m[1] if m[1]
    PLUG_OPTS[:blit_port] = m[2].to_i
  end

  opts.on('--[no-]target-tls', 'enable/disable TLS to target') { |t| PLUG_OPTS[:tgt_tls] = t }
  opts.on('--[no-]server-tls', 'enable/disable TLS to clients') { |t| PLUG_OPTS[:svr_tls] = t }
end

begin
  opts.parse!(ARGV)
rescue StandardError
  (warn $!
   exit 1)
end

# Get target/listen argument
rx = /^([\w.]+):(\d+)(?:@(?:([\w.]+):)?(\d+))?$/
unless (m = rx.match(ARGV.shift)) and ARGV.shift.nil?
  warn opts.banner
  exit 1
end

PLUG_OPTS[:tgtaddr] = m[1]
PLUG_OPTS[:tgtport] = m[2].to_i
PLUG_OPTS[:svraddr] ||= m[3] || '0.0.0.0'
PLUG_OPTS[:svrport] ||= (m[4] || PLUG_OPTS[:tgtport]).to_i

EventMachine.run do
  # Instantiate controller
  BlitPlug::Controller.new(PLUG_OPTS[:tgtaddr], PLUG_OPTS[:tgtport], BlitPlug::TCPClient)

  # Start event loop
  BlitPlug::UI.log "%Starting TCP PlugServer #{PLUG_OPTS[:svraddr]}:#{PLUG_OPTS[:svrport]} -> #{PLUG_OPTS[:tgtaddr]}:#{PLUG_OPTS[:tgtport]}"

  EventMachine.start_server(PLUG_OPTS[:svraddr], PLUG_OPTS[:svrport], BlitPlug::TCPListener)
end
