#!/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 = {}

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

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

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

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

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

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

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

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

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

parser.on("--source-port PORT (optional)") do |x|
	opts[:source_port] = x
end

parser.on("--custom-nmap-options OPTIONS (optional)") do |x|
	opts[:custom_nmap_options] = x
end

parser.on("--disable-udp-probes (optional)") do
	opts[:disable_udp_probes] = true
end

parser.on("--disable-finger-users (optional)") do
	opts[:disable_finger_users] = true
end

parser.on("--disable-snmp-scan (optional)") do 
	opts[:disable_snmp_scan] = true
end

parser.on("--disable-service-identification (optional)") do
	opts[:disable_service_identification] = true
end

parser.on("--smb-user USER (optional)") do |x|
	opts[:smb_user] = x
end

parser.on("--smb-pass PASS (optional)") do |x|
	opts[:smb_pass] = x
end

parser.on("--smb-domain DOMAIN (optional)") do |x|
	opts[:smb_domain] = x
end

parser.on("--dry-run (optional)") do
	opts[:dry_run] = true
end

parser.on("--single-scan (optional)") do
	opts[:single_scan] = true
end

parser.on("--fast-detect (optional)") do
	opts[:fast_detect] = true
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

# Provide default values for certain options - If there's no alternative set
# use the default provided by Pro -- see the documentation.
project 			= opts[:project]	|| usage(parser)
targets 			= opts[:targets]	|| usage(parser)
blacklist			= opts[:blacklist]
speed				= opts[:speed]		|| "Insane"
extra_ports			= opts[:extra_ports]
blacklist_ports			= opts[:blacklist_ports]
custom_ports			= opts[:custom_ports]
portscan_timeout		= opts[:portscan_timeout]	|| 300
source_port			= opts[:source_port]
custom_nmap_options		= opts[:custom_nmap_options] || 
disable_udp_probes		= opts[:disable_udp_probes] || false
disable_finger_users		= opts[:disable_finger_users] || false
disable_snmp_scan		= opts[:disable_snmp_scan] || false
disable_service_identification	= opts[:disable_service_identification] || false
smb_user			= opts[:smb_user] || ""
smb_pass			= opts[:smb_pass] || ""
smb_domain			= opts[:smb_domain] || ""
single_scan			= opts[:single_scan] || false
fast_detect			= opts[:fast_detect] || false

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

# Create the task object with all options
task 		= @rpc.call("pro.start_discover", {
				'workspace'		=> project,
				'username' 		=> user,
				'ips'			=> targets,
				'DS_BLACKLIST_HOSTS'	=> blacklist,
				'DS_PORTSCAN_SPEED'	=> speed,
				'DS_PORTS_EXTRA'	=> extra_ports,
				'DS_PORTS_BLACKLIST'	=> blacklist_ports,
				'DS_PORTS_CUSTOM'	=> custom_ports,
				'DS_PORTSCAN_TIMEOUT' 	=> portscan_timeout,
				'DS_PORTSCAN_SOURCE_PORT' => source_port,
				'DS_CustomNmap'		=> custom_nmap_options,
				'DS_UDP_PROBES'		=> disable_udp_probes,
				'DS_FINGER_USERS'	=> disable_finger_users,
				'DS_SNMP_SCAN'		=> disable_snmp_scan,
				'DS_IDENTIFY_SERVICES'	=> disable_service_identification,
				'DS_SMBUser'		=> smb_user,
				'DS_SMBPass'		=> smb_pass,
				'DS_SMBDomain'		=> smb_domain,
				'DS_SINGLE_SCAN'	=> single_scan, 
				'DS_FAST_DETECT'	=> fast_detect
})

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!"