#
# sf - Sept 2010
#
require 'thread'
require 'rex/logging'
require 'rex/socket'

module Rex
module Proto
module Proxy

#
# A Socks4a proxy server.
#
class Socks4a

	#
	# A client connected to the Socks4a server.
	#
	class Client

		REQUEST_VERSION                 = 4
		REPLY_VERSION                   = 0

		COMMAND_CONNECT                 = 1
		COMMAND_BIND                    = 2

		REQUEST_GRANTED                 = 90
		REQUEST_REJECT_FAILED           = 91
		REQUEST_REJECT_CONNECT          = 92
		REQUEST_REJECT_USERID           = 93

		HOST                            = 1
		PORT                            = 2

		#
		# A Socks4a packet.
		#
		class Packet

			def initialize
				@version   = REQUEST_VERSION
				@command   = 0
				@dest_port = 0
				@dest_ip   = '0.0.0.0'
				@userid    = ''
			end

			#
			# A helper function to recv in a Socks4a packet byte by byte.
			#
			# sf: we could just call raw = sock.get_once but some clients
			#     seem to need reading this byte by byte instead.
			#
			def Packet.recv( sock, timeout=30 )
				raw = ''
				# read in the 8 byte header
				while( raw.length < 8 )
					raw << sock.read( 1 )
				end
				# if its a request there will be more data
				if( raw[0..0].unpack( 'C' ).first == REQUEST_VERSION )
					# read in the userid
					while( raw[8..raw.length].index( "\x00" ) == nil )
						raw << sock.read( 1 )
					end
					# if a hostname is going to be present, read it in
					ip = raw[4..7].unpack( 'N' ).first
					if( ( ip & 0xFFFFFF00 ) == 0x00000000 and ( ip & 0x000000FF ) != 0x00 )
						hostname = ''
						while( hostname.index( "\x00" ) == nil )
							hostname += sock.read( 1 )
						end
						raw << hostname
					end
				end
				# create a packet from this raw data...
				packet = Packet.new
				packet.from_r( raw ) ? packet : nil
			end

			#
			# Pack a packet into raw bytes for transmitting on the wire.
			#
			def to_r
				raw = [ @version, @command, @dest_port, Rex::Socket.addr_atoi( @dest_ip ) ].pack( 'CCnN' )
				return raw if( @userid.empty? )
				return raw + [ @userid ].pack( 'Z*' )
			end

			#
			# Unpack a raw packet into its components.
			#
			def from_r( raw )
				return false if( raw.length < 8 )
				@version   = raw[0..0].unpack( 'C' ).first
				return false if( @version != REQUEST_VERSION and @version != REPLY_VERSION )
				@command   = raw[1..1].unpack( 'C' ).first
				@dest_port = raw[2..3].unpack( 'n' ).first
				@dest_ip   = Rex::Socket.addr_itoa( raw[4..7].unpack( 'N' ).first )
				if( raw.length > 8 )
					@userid = raw[8..raw.length].unpack( 'Z*' ).first
					# if this is a socks4a request we can resolve the provided hostname
					if( self.is_hostname? )
						hostname = raw[(8+@userid.length+1)..raw.length].unpack( 'Z*' ).first
						@dest_ip = self.resolve( hostname )
						# fail if we couldnt resolve the hostname
						return false if( not @dest_ip )
					end
				else
					@userid  = ''
				end
				return true
			end

			def is_connect?
				@command == COMMAND_CONNECT ? true : false
			end

			def is_bind?
				@command == COMMAND_BIND ? true : false
			end

			attr_accessor :version, :command, :dest_port, :dest_ip, :userid

			protected

			#
			# Resolve the given hostname into a dotted IP address.
			#
			def resolve( hostname )
				if( not hostname.empty? )
					begin
						return Rex::Socket.addr_itoa( Rex::Socket.gethostbyname( hostname )[3].unpack( 'N' ).first )
					rescue ::SocketError
						return nil
					end
				end
				return nil
			end

			#
			# As per the Socks4a spec, check to see if the provided dest_ip is 0.0.0.XX
			# which indicates after the @userid field contains a hostname to resolve.
			#
			def is_hostname?
				ip = Rex::Socket.addr_atoi( @dest_ip )
				if( ip & 0xFFFFFF00 == 0x00000000 )
					return true if( ip & 0x000000FF != 0x00 )
				end
				return false
			end

		end

		#
		# A mixin for a socket to perform a relay to another socket.
		#
		module Relay

			#
			# Relay data coming in from relay_sock to this socket.
			#
			def relay( relay_client, relay_sock )
				@relay_client = relay_client
				@relay_sock   = relay_sock
				# start the relay thread (modified from Rex::IO::StreamAbstraction)
				@relay_thread = Rex::ThreadFactory.spawn("SOCKS4AProxyServerRelay", false) do
					loop do
						closed = false
						buf    = nil

						begin
							s = Rex::ThreadSafe.select( [ @relay_sock ], nil, nil, 0.2 )
							if( s == nil || s[0] == nil )
								next
							end
						rescue
							closed = true
						end

						if( closed == false )
							begin
								buf = @relay_sock.sysread( 32768 )
								closed = true if( buf == nil )
							rescue
								closed = true
							end
						end

						if( closed == false )
							total_sent   = 0
							total_length = buf.length
							while( total_sent < total_length )
								begin
									data = buf[total_sent, buf.length]
									sent = self.write( data )
									if( sent > 0 )
										total_sent += sent
									end
								rescue
									closed = true
									break
								end
							end
						end

						if( closed )
							@relay_client.stop
							::Thread.exit
						end
					end
				end

			end

		end

		#
		# Create a new client connected to the server.
		#
		def initialize( server, sock )
			@server        = server
			@lsock         = sock
			@rsock         = nil
			@client_thread = nil
			@mutex         = ::Mutex.new
		end

		#
		# Start handling the client connection.
		#
		def start
			# create a thread to handle this client request so as to not block the socks4a server
			@client_thread = Rex::ThreadFactory.spawn("SOCKS4AProxyClient", false) do
				begin
					@server.add_client( self )
					# get the initial client request packet
					request = Packet.recv( @lsock )
					raise "Invalid Socks4 request packet received." if not request
					# handle the request
					begin
						# handle socks4a conenct requests
						if( request.is_connect? )
							# perform the connection request
							params = {
								'PeerHost' => request.dest_ip,
								'PeerPort' => request.dest_port,
							}
							params['Context'] = @server.opts['Context'] if @server.opts.has_key?('Context')

							@rsock = Rex::Socket::Tcp.create( params )
							# and send back success to the client
							response         = Packet.new
							response.version = REPLY_VERSION
							response.command = REQUEST_GRANTED
							@lsock.put( response.to_r )
						# handle socks4a bind requests
						elsif( request.is_bind? )
							# create a server socket for this request
							params = {
								'LocalHost' => '0.0.0.0',
								'LocalPort' => 0,
							}
							params['Context'] = @server.opts['Context'] if @server.opts.has_key?('Context')
							bsock = Rex::Socket::TcpServer.create( params )
							# send back the bind success to the client
							response           = Packet.new
							response.version   = REPLY_VERSION
							response.command   = REQUEST_GRANTED
							response.dest_ip   = '0.0.0.0'
							response.dest_port = bsock.getlocalname()[PORT]
							@lsock.put( response.to_r )
							# accept a client connection (2 minute timeout as per spec)
							begin
								::Timeout.timeout( 120 ) do
									@rsock = bsock.accept
								end
							rescue ::Timeout::Error
								raise "Timeout reached on accept request."
							end
							# close the listening socket
							bsock.close
							# verify the connection is from the dest_ip origionally specified by the client
							rpeer = @rsock.getpeername
							raise "Got connection from an invalid peer." if( rpeer[HOST] != request.dest_ip )
							# send back the client connect success to the client
							#
							# sf: according to the spec we send this response back to the client, however
							#     I have seen some clients who bawk if they get this second response.
							#
							response           = Packet.new
							response.version   = REPLY_VERSION
							response.command   = REQUEST_GRANTED
							response.dest_ip   = rpeer[HOST]
							response.dest_port = rpeer[PORT]
							@lsock.put( response.to_r )
						else
							raise "Unknown request command received #{request.command} received."
						end
					rescue
						# send back failure to the client
						response         = Packet.new
						response.version = REPLY_VERSION
						response.command = REQUEST_REJECT_FAILED
						@lsock.put( response.to_r )
						# raise an exception to close this client connection
						raise "Failed to handle the clients request."
					end
					# setup the two way relay for full duplex io
					@lsock.extend( Relay )
					@rsock.extend( Relay )
					# start the socket relays...
					@lsock.relay( self, @rsock )
					@rsock.relay( self, @lsock )
				rescue
					wlog( "Client.start - #{$!}" )
					self.stop
				end
			end
		end

		#
		# Stop handling the client connection.
		#
		def stop
			@mutex.synchronize do
				if( not @closed )

					begin
						@lsock.close if @lsock
					rescue
					end

					begin
						@rsock.close if @rsock
					rescue
					end

					@client_thread.kill if( @client_thread and @client_thread.alive? )

					@server.remove_client( self )

					@closed = true
				end
			end
		end

	end

	#
	# Create a new Socks4a server.
	#
	def initialize( opts={} )
		@opts          = { 'ServerHost' => '0.0.0.0', 'ServerPort' => 1080 }
		@opts          = @opts.merge( opts )
		@server        = nil
		@clients       = ::Array.new
		@running       = false
		@server_thread = nil
	end

	#
	# Check if the server is running.
	#
	def is_running?
		return @running
	end

	#
	# Start the Socks4a server.
	#
	def start
			begin
				# create the servers main socket (ignore the context here because we don't want a remote bind)
				@server = Rex::Socket::TcpServer.create( 'LocalHost' => @opts['ServerHost'], 'LocalPort' => @opts['ServerPort'] )
				# signal we are now running
				@running = true
				# start the servers main thread to pick up new clients
				@server_thread = Rex::ThreadFactory.spawn("SOCKS4AProxyServer", false) do
					while( @running ) do
						begin
							# accept the client connection
							sock = @server.accept
							# and fire off a new client instance to handle it
							Client.new( self, sock ).start
						rescue
							wlog( "Socks4a.start - server_thread - #{$!}" )
						end
					end
				end
			rescue
				wlog( "Socks4a.start - #{$!}" )
				return false
			end
			return true
	end

	#
	# Block while the server is running.
	#
	def join
		@server_thread.join if @server_thread
	end

	#
	# Stop the Socks4a server.
	#
	def stop
		if( @running )
			# signal we are no longer running
			@running = false
			# stop any clients we have (create a new client array as client.stop will delete from @clients)
			clients = []
			clients.concat( @clients )
			clients.each do | client |
				client.stop
			end
			# close the server socket
			@server.close if @server
			# if the server thread did not terminate gracefully, kill it.
			@server_thread.kill if( @server_thread and @server_thread.alive? )
		end
		return !@running
	end

	def add_client( client )
		@clients << client
	end

	def remove_client( client )
		@clients.delete( client )
	end

	attr_reader :opts

end

end; end; end