2013-07-29

Bouncer aggiornado

Después de mucho tiempo volví a programar un toque en ruby. Por necesidades laborales, tuve que usar de nuevo mi bouncer.rb, pero lamentablemente "rdoc/usage" no está en ruby 1.9. Así que lo adapté para que no lo necesite más. Abajo pongo el código:

#!/usr/bin/env ruby

# == Synopsis
#
# redirect.rb: Redirects TCP connections to distant machines. Handles simultaneously many connections.
#
# == Usage
#
# ruby redirect.rb [OPTION]
#
# -h, --help:
#   Show help
#
# --ip ip, -i ip:
#   Accept connections from ip (default 127.0.0.1)
#
# --port port, -p port:
#   Listen on port (default 12345)
#
# --target ip:port, -t ip:port
#   Connect to ip:port (default 127.0.0.1:23456)
#

require 'getoptlong'
#require 'rdoc/usage'
require 'socket'

class Redirect
 def initialize( source_ip, listening_port, target_ip, target_port )
  @source_ip = source_ip
  @target_ip = target_ip
  @target_port = target_port

  @server_socket = TCPServer.new( "", listening_port )
    @server_socket.setsockopt( Socket::SOL_SOCKET, Socket::SO_REUSEADDR, 1 )
   
  @descriptors = [ @server_socket ]
  @next_step = {}
 end

 def handle_new_connection
  incoming = @server_socket.accept
  if incoming.peeraddr[3] == @source_ip
   begin
    outgoing = TCPSocket.new( @target_ip, @target_port )
   
    @next_step[ incoming ] = outgoing
    @next_step[ outgoing ] = incoming

    @descriptors += [ incoming, outgoing ]
   rescue
    silent_close( incoming )
   end
  else
   silent_close( incoming )
  end
 end

 def silent_close( sock )
  begin
   sock.close
  rescue
   #do nothing intentionally
  end
 end

  def propagate(sock)
    next_sock = @next_step[sock]
    next_sock.write(sock.read_nonblock(1000 * 1000))
  end

 def finish_connection(sock)
  next_sock = @next_step[sock]
  [ sock, next_sock ].each do
   |s|
   silent_close(s)
   @descriptors.delete(s)
   @next_step.delete(s)
  end
 end

 def run
  loop do
   connections = select( @descriptors )
   connections[0].each do
    |sock|
    if sock == @server_socket then
     handle_new_connection
    else
     begin
      sock.eof? ? finish_connection(sock) : propagate(sock)
     rescue
      finish_connection(sock)
     end
    end
   end
  end
 end
end

if $0 == __FILE__ then

 opts = GetoptLong.new(
   [ '--help', '-h', GetoptLong::NO_ARGUMENT ],
   [ '--ip', '-i', GetoptLong::REQUIRED_ARGUMENT ],
   [ '--port', '-p', GetoptLong::REQUIRED_ARGUMENT ],
   [ '--target', '-t', GetoptLong::REQUIRED_ARGUMENT ]
 )

 ip = '127.0.0.1'
 port = '12345'
 target = '127.0.0.1:23456'

 opts.each do
  |opt, arg|
  case opt
  when '--help'
   print """\
Usage:
ruby bouncer.rb [options]

options:
 -i [ip] --ip=[ip]
 -p [port] --port=[port]
 -t [host:port] --target=[host:port]
"""
   exit
  when '--ip'
   ip = arg
  when '--port'
   port = arg
  when '--target'
   target = arg
  end
 end

 port = port.to_i
 target = target.split(":")

 trap("SIGINT") do
   exit
 end

 Redirect.new(ip, port, target[0], target[1].to_i).run
end



Happy hacking,
Aureliano