# -*- coding: binary -*-
#
# sf - Sept 2010
#
require 'thread'
require 'rex/logging'
require 'rex/socket'

module Rex
module Proto
module Proxy

#
# A Socks4a proxy server.
#
class Socks5

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



    REPLY_VERSION                   = 0
    REQUEST_VERSION                 = 5

    ATYP_IP4                        = 1
    ATYP_DOMAINNAME                 = 3
    ATYP_IP6                        = 4

    AUTH_NOT_REQUIRED               = 0
    AUTH_GSSAPI                     = 1
    AUTH_USERNAME_PASSWORD          = 2
    AUTH_NO_ACCEPTABLE_METHODS      = 255

    COMMAND_CONNECT                 = 1
    COMMAND_BIND                    = 2
    COMMAND_UDP_ASSOCIATE           = 3

    REQUEST_GRANTED                 = 0
    REQUEST_REJECT_GENERAL          = 1

    HOST                            = 1
    PORT                            = 2

    #
    # A Socks4a packet.
    #
    class Packet

      def initialize
        @version    = REQUEST_VERSION
        @command    = 0
        @dest_port  = 0
        @dest_host  = '0.0.0.0'
        @dest_atype = 0
      end

      #
      # A helper function to recv in a Socks5 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 4 byte header
        while raw.length < 4
          raw << sock.read(1)
        end

        raw = raw.unpack("CCCC")
        return nil unless raw[0] == REQUEST_VERSION
        return nil unless [COMMAND_CONNECT, COMMAND_BIND, COMMAND_UDP_ASSOCIATE].include? raw[1]
        return nil unless raw[2] == 0
        return nil unless [ATYP_IP4, ATYP_DOMAINNAME, ATYP_IP6].include? raw[3]

        # read in address
        case raw[3]
        when ATYP_IP4
          address_length = 4
        when ATYP_DOMAINNAME
          address_length = sock.read(1).unpack("C").first
        when ATYP_IP6
          address_length = 16
        end
        address = sock.read(address_length)
        return nil unless address.to_s.length == address_length
        if raw[3] != ATYP_DOMAINNAME
          address = Rex::Socket.addr_ntoa(address)
        end
        # read in 2 bytes ALSO have unpack as a 16-bit integer

        dest_port = sock.read(2).unpack("n")
        pkt = {
          :version   => raw[0],
          :command   => raw[1],
          :dest_port => dest_port,
          :dest_host => address,
          :dest_atype => raw[3]
        }

        # create a packet from this raw data...
        packet = Packet.new
        # packet.from_h(pkt) ? packet : nil
        packet.from_h(pkt) ? packet : nil
      end

      #
      # Pack a packet into raw bytes for transmitting on the wire.
      #
      def to_r
        raw = [ @version, @command, 0, @dest_atype ].pack('CCCC')

        if is_hostname?
          raw << [@dest_host.length].pack('C')
        end
        raw << @dest_host
        raw << [@dest_port].pack('n')
        raw
      end

      #
      # Unpack a raw packet into its components.
      #
      def from_h(pkt)
        @version    = pkt[:version]
        @command    = pkt[:command]
        @dest_port  = pkt[:dest_port]
        @dest_host  = pkt[:dest_host]
        @dest_atype = pkt[:dest_atype]
        # if this is a socks4a request we can resolve the provided hostname
        if self.is_hostname?
          @dest_host = self.resolve(pkt[:dest_host])
          # fail if we couldnt resolve the hostname
          return false unless @dest_host
        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_host, :dest_atype

      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?
        return @dest_atype == ATYP_DOMAINNAME
      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("SOCKS5ProxyServerRelay", 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

    #
    # Checks if the request supports Authentication methods.
    #
    def authenticate
      # reads two bytes of the request
      request = @lsock.read(2)
      # checks if the request is long enough
      return false unless request.to_s.length == 2
      # unpacks request from string to an array of two integers
      request = request.unpack('CC')
      # checks if the client in running the same version
      return false unless request[0] == REQUEST_VERSION
      # reads the rest of the request
      auth_methods = @lsock.read(request[1])
      # checks if auth_methods got the same amount as in the request
      return false unless auth_methods.to_s.length == request[1]
      # checks that there is something in the request
      return false unless auth_methods.to_s.length > 0
      # this unpacks everything from auth_methods instead of just 2
      auth_methods = auth_methods.unpack('C*')
      # checks if auth is required
      if auth_methods.include? AUTH_NOT_REQUIRED
        response = [REQUEST_VERSION, AUTH_NOT_REQUIRED].pack('CC')
        @lsock.write(response)
        wlog("-Socks5 Auth- Client was authenticated Auth was not required.")
        # returns true because the client was authenticated
        return true
      else
        response = [REQUEST_VERSION, AUTH_NO_ACCEPTABLE_METHODS].pack('CC')
        @lsock.write(response)
        wlog("-Sock5 Auth- Client was not authenicated because no acceptable methods were used.")
      end
      false
    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("SOCKS5ProxyClient", false) do
        begin
          @server.add_client( self )
          unless authenticate
            wlog("-Auth- Authentication Failed")
            self.stop
            return
          end
          # get the initial client request packet
          request = Packet.recv( @lsock )
          raise "Invalid Socks5 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_host,
                'PeerPort' => request.dest_port,
              }
              params['Context'] = @server.opts['Context'] if @server.opts.has_key?('Context')
              wlog('Marker 1')
              @rsock = Rex::Socket::Tcp.create(params)
              # and send back success to the client
              response             = Packet.new
              response.version     = REPLY_VERSION
              response.command     = REQUEST_GRANTED
              response.dest_port   = request.dest_port
              response.dest_host   = request.dest_host
              response.dest_atype   = request.dest_atype
              wlog('Marker 2')
              @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 )
              wlog('Marker 1')
            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_GENERAL
            @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("SOCKS5ProxyServer", 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( "Socks5.start - server_thread - #{$!}" )
            end
          end
        end
      rescue
        wlog( "Socks5.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