# == 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'
RDoc::usage
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