#!/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 )
      unless PLUG_OPTS[:quiet]
        PLUG_OPTS[:out].puts msg
      end
    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 = Array.new

      ## 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 = Array.new
      @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)
  STDERR.puts 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] = 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

  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)
      STDERR.puts "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

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


EventMachine::run {
  # Instantiate controller
  ctrl = 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)
}