#!/usr/bin/env ruby
#
#

require 'rubygems'
require 'eventmachine'
require 'socket'
require 'optparse'
require 'rbkb'

def bail(*msg)
  STDERR.puts msg
  exit 1
end

class Plug
  module UI
    def log( *msg )
      unless PLUG_OPTS[:quiet]
        PLUG_OPTS[:out].puts msg
      end
    end
    module_function :log
  end

  class Controller
    attr_accessor :tgtaddr, :tgtport, :tgtclient
    @@controllers=nil

    def initialize(tgtaddr, tgtport, tgtclient)

      @tgtaddr = tgtaddr
      @tgtport = tgtport
      @tgtclient = tgtclient

      @@controllers = self
    end

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

    def dispatch_rcv(snder, data)
      data # XXX for now
    end

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

    def dispatch_close(snder)
      nil # XXX for now
    end

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

    def self.proxy(cli)
      unless (ctrl = @@controllers)
        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

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


  module BaseTCP
    include UI

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

    def post_init
      @plug_peers = Array.new
      @kind = :conn  # default
    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
      log "%#{kind.to_s.upcase}-#{@sock_peername}-CONNECTED"
    end


    def unbind
      log "%#{kind.to_s.upcase}-#{@sock_peername}-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 Plug::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 = Plug::Controller.proxy(self)

      notify_connection
    end

  end # module TCPListener


  module TCPClient
    include Plug::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
    end

  end # module TCPClient

end # module Plug

PLUG_OPTS={ :quiet => false, :out => STDOUT }

  if __FILE__ == $0


  #############################################################################
  ### MAIN
  #############################################################################
  #
  # Get arguments
  opts = OptionParser.new do |opts|
    opts.banner = "Usage: #{$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] = File.open(o, "w") rescue (bail $!)
    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)
        STDERR.puts "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

  end

  opts.parse!(ARGV) rescue (STDERR.puts $!; exit 1)


  # Get target/listen argument
  rx = /^([\w\.]+):(\d+)(?:@(?:([\w\.]+):)?(\d+))?$/
  unless (m = rx.match(ARGV.shift)) and ARGV.shift.nil?
    $stderr.puts 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


  # Start controller
  ctrl = Plug::Controller.new(PLUG_OPTS[:tgtaddr], PLUG_OPTS[:tgtport], Plug::TCPClient)

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

  EventMachine::run {
    EventMachine::start_server(PLUG_OPTS[:svraddr], PLUG_OPTS[:svrport], Plug::TCPListener)
  }


end