##
# $Id: $
##

##
#
# RFB protocol support
#
# by Joshua J. Drake <jduck>
#
# Based on:
# vnc_auth_none contributed by Matteo Cantoni <goony[at]nothink.org>
# vnc_auth_login contributed by carstein <carstein.sec[at]gmail.com>
#
# TODO: determine how to detect a view-only session.
##

module Rex
module Proto
module RFB

class Client

	def initialize(sock, opts = {})
		@sock = sock
		@opts = opts

		@banner = nil
		@majver = MajorVersion
		@minver = -1
		@auth_types = []
	end

	def read_error_message
		len = @sock.get_once(4)
		return 'Unknown error' if not len or len.length != 4

		len = len.unpack("N").first
		@sock.get_once(len)
	end

	def handshake
		@banner = @sock.get_once(12)
		if not @banner
			@error = "Unable to obtain banner from server"
			return false
		end

		# RFB Protocol Version 3.3 (1998-01)
		# RFB Protocol Version 3.7 (2003-08)
		# RFB Protocol Version 3.8 (2007-06)

		if @banner =~ /RFB ([0-9]{3})\.([0-9]{3})/
			maj = $1.to_i
			if maj != MajorVersion
				@error = "Invalid major version number: #{maj}"
				return false
			end
		else
			@error = "Invalid RFB banner: #{@banner}"
			return false
		end

		@minver = $2.to_i

		our_ver = "RFB %03d.%03d\n" % [MajorVersion, @minver]
		@sock.put(our_ver)

		true
	end

	def connect(password = nil)
		return false if not handshake
		return false if not authenticate(password)
		return false if not send_client_init
		true
	end

	def send_client_init
		if @opts[:exclusive]
			@sock.put("\x00") # do share.
		else
			@sock.put("\x01") # do share.
		end
	end

	def authenticate(password = nil)
		type = negotiate_authentication
		return false if not type

		# Authenticate.
		case type
		when AuthType::None
			# Nothing here.

		when AuthType::VNC
			return false if not negotiate_vnc_auth(password)

		end

		# Handle reading the security result message
		result = @sock.get_once(4)
		if not result
			@error = "Unable to read auth result"
			return false
		end

		result = result.unpack('N').first
		case result
		when 0
			return true

		when 1
			if @minver >= 8
				msg = read_error_message
				@error = "Authentication failed: #{msg}"
			else
				@error = "Authentication failed"
			end
		when 2
			@error = "Too many authentication attempts"
		else
			@error = "Unknown authentication result: #{result}"
		end

		false
	end

	def negotiate_authentication
		# Authentication type negotiation is protocol version specific.
		if @minver < 7
			buf = @sock.get_once(4)
			if not buf
				@error = "Unable to obtain requested authentication method"
				return nil
			end
			@auth_types = buf.unpack('N')
			if not @auth_types or @auth_types.first == 0
				msg = read_error_message
				@error = "No authentication types available: #{msg}"
				return nil
			end
		else
			buf = @sock.get_once(1)
			if not buf
				@error = "Unable to obtain supported authentication method count"
				return nil
			end

			# first byte is number of security types
			num_types = buf.unpack("C").first
			if (num_types == 0)
				msg = read_error_message
				@error = "No authentication types available: #{msg}"
				return nil
			end

			buf = @sock.get_once(num_types)
			if not buf or buf.length != num_types
				@error = "Unable to read authentication types"
				return nil
			end

			@auth_types = buf.unpack("C*")
		end

		if not @auth_types or @auth_types.length < 1
			@error = "No authentication types found"
			return nil
		end

		# Select the one we prefer
		selected = nil
		selected ||= AuthType::None if @opts[:allow_none] and @auth_types.include? AuthType::None
		selected ||= AuthType::VNC if @auth_types.include? AuthType::VNC

		if not selected
			@error = "No supported authentication method found."
			return nil
		end

		# For 3.7 and later, clients must state which security-type to use
		@sock.put([selected].pack('C')) if @minver >= 7

		selected
	end

	def negotiate_vnc_auth(password = nil)
		challenge = @sock.get_once(16)
		if not challenge or challenge.length != 16
			@error = "Unable to obtain VNC challenge"
			return false
		end

		response = Cipher.encrypt(challenge, password)
		@sock.put(response)

		true
	end

	attr_reader :error, :majver, :minver, :auth_types
	attr_reader :view_only
end

end
end
end