#!/usr/bin/env ruby
require 'rubygems'
require 'optparse'
require 'msfrpc-client'
require 'rex/ui'

def usage(ropts)
	$stderr.puts ropts

	if @rpc and @rpc.token
		wspaces = @rpc.call("pro.workspaces") rescue {}
		if wspaces.keys.length > 0
			$stderr.puts "Active Projects:"
			wspaces.each_pair do |k,v|
				$stderr.puts "\t#{k}"
			end
		end
	end
	$stderr.puts ""
	exit(1)
end

opts = {}
opts[:blacklist]      = ''
opts[:whitelist_ports] = ''
opts[:blacklist_ports] = ''
opts[:exploit_timeout] = 5
opts[:limit_sessions] = true
opts[:ignore_fragile_devices] = true
opts[:filter_by_os]   = true
opts[:only_match]     = false
opts[:match_vulns]    = true
opts[:match_ports]    = true
opts[:payload_method] = "auto"
opts[:payload_type]   = "meterpreter"
opts[:payload_ports]  = "4000-5000"
opts[:evasion_level_tcp] = 0
opts[:evasion_level_app] = 0
opts[:module_filter] = ''

# Parse script-specific options
parser = Msf::RPC::Client.option_parser(opts)
parser.separator('Exploit Specific Options:')

parser.on("--project PROJECT") do |x|
	opts[:project] = x
end

parser.on("--targets TARGETS") do |x|
	opts[:targets] = x
end

parser.on("--speed SPEED") do |x|
	opts[:speed] = x
end

parser.on("--minimum-rank RANK") do |x|
	opts[:rank] = x
end

parser.on("--blacklist BLACKLIST (optional)") do |x|
	opts[:blacklist] = x
end

parser.on("--whitelist-ports PORTS (optional)") do |x|
	opts[:whitelist_ports] = x
end

parser.on("--blacklist-ports PORTS (optional)") do |x|
	opts[:blacklist_ports] = x
end

parser.on("--exploit-timeout TIMEOUT (optional)") do |x|
	opts[:exploit_timeout] = x
end

parser.on("--limit-sessions (optional)") do |x|
	opts[:limit_sessions] = (x =~ /^(y|t|1)/i ? true : false )
end

parser.on("--ignore-fragile-devices (optional)") do |x|
	opts[:ignore_fragile_devices] = (x =~ /^(y|t|1)/i ? true : false )
end

parser.on("--filter-by-os (optional)") do |x|
	opts[:filter_by_os] = (x =~ /^(y|t|1)/i ? true : false )
end

parser.on("--dry-run (optional)") do |x|
	opts[:only_match] = (x =~ /^(y|t|1)/i ? true : false )
end

parser.on("--match-vulns (optional)") do |x|
	opts[:match_vulns] = (x =~ /^(y|t|1)/i ? true : false )
end

parser.on("--match-ports (optional)") do |x|
	opts[:match_ports] = (x =~ /^(y|t|1)/i ? true : false )
end

parser.on("--payload-method AUTO|REVERSE|BIND (optional)") do |x|
	opts[:payload_method] = x
end

parser.on("--payload-type METERPRETER|SHELL (optional)") do |x|
	opts[:payload_type] = x
end

parser.on("--payload-ports PORTS (optional)") do |x|
	opts[:payload_ports] = x
end

parser.on("--evasion-level-tcp LEVEL (optional)") do |x|
	opts[:evasion_level_tcp] = x
end

parser.on("--evasion-level-app LEVEL (optional)") do |x|
	opts[:evasion_level_app] = x
end

parser.on("--module-filter FILTER (optional)") do |x|
	opts[:module_filter] = x
end

parser.on("--help") do
	$stderr.puts parser
	exit(1)
end

parser.separator('')
parser.parse!(ARGV)

@rpc  = Msf::RPC::Client.new(opts)

if not @rpc.token
	$stderr.puts "Error: Invalid RPC server options specified"
	$stderr.puts parser
	exit(1)
end

# Store the user's settings
project 			= opts[:project]	|| usage(parser)
targets 			= opts[:targets]	|| usage(parser)
rank				= opts[:rank]		|| usage(parser)
speed				= opts[:speed]		|| usage(parser)
blacklist			= opts[:blacklist]
whitelist_ports			= opts[:whitelist_ports]
blacklist_ports			= opts[:blacklist_ports]
exploit_timeout			= opts[:exploit_timeout] 
limit_sessions			= opts[:limit_sessions]
ignore_fragile_devices		= opts[:ignore_fragile_devices] 
filter_by_os			= opts[:filter_by_os]
only_match			= opts[:only_match]
match_vulns			= opts[:match_vulns]
match_ports			= opts[:match_ports]
payload_method			= opts[:payload_method]
payload_type			= opts[:payload_type]
payload_ports			= opts[:payload_ports]
evasion_level_tcp		= opts[:evasion_level_tcp]
evasion_level_app		= opts[:evasion_level_app]
module_filter			= opts[:module_filter]
#===

# Get the default user
user   		= @rpc.call("pro.default_admin_user")['username']

# Create the task object with all options
task 		= @rpc.call("pro.start_exploit", {
				'workspace'			=> project,
				'username' 			=> user,
				'DS_WHITELIST_HOSTS'		=> targets,
				'DS_BLACKLIST_HOSTS'		=> blacklist,
				'DS_WHITELIST_PORTS'    	=> whitelist_ports,
				'DS_BLACKLIST_PORTS'		=> blacklist_ports, 
				'DS_MinimumRank' 		=> rank, 
				'DS_EXPLOIT_SPEED'		=> speed, 
				'DS_EXPLOIT_TIMEOUT'		=> exploit_timeout,
				'DS_LimitSessions'		=> limit_sessions,
				'DS_IgnoreFragileDevices' 	=> ignore_fragile_devices, 
				'DS_FilterByOS'			=> filter_by_os, 
				'DS_OnlyMatch'			=> only_match,
				'DS_MATCH_VULNS'		=> match_vulns,
				'DS_MATCH_PORTS'		=> match_ports, 
				'DS_PAYLOAD_METHOD'		=> payload_method, 
				'DS_PAYLOAD_TYPE'		=> payload_type, 
				'DS_PAYLOAD_PORTS'		=> payload_ports, 
				'DS_EVASION_LEVEL_TCP'		=> evasion_level_tcp, 
				'DS_EVASION_LEVEL_APP'		=> evasion_level_app,
				'DS_ModuleFilter'		=> module_filter
})

puts "DEBUG: Running task with #{task.inspect}"

if not task['task_id']
	$stderr.puts "[-] Error starting the task: #{task.inspect}"
	exit(0)
end

puts "[*] Creating Task ID #{task['task_id']}..."
while true
	select(nil, nil, nil, 0.50)

	stat = @rpc.call("pro.task_status", task['task_id'])

	if stat['status'] == 'invalid'
		$stderr.puts "[-] Error checking task status"
		exit(0)
	end

	info = stat[ task['task_id'] ]

	if not info
		$stderr.puts "[-] Error finding the task"
		exit(0)
	end

	if info['status'] == "error"
		$stderr.puts "[-] Error generating report: #{info['error']}"
		exit(0)
	end

	break if info['progress'] == 100
end

$stdout.puts "[+] Task Complete!"
