
#   Copyright (C) Stonesoft Corporation 2010 - 2013.
#   All rights reserved.
#
#   The StoneGate software, manuals, and technical
#   literature may not be reproduced in any form or
#   by any means except by permission in writing from
#   Stonesoft Corporation.
#    ___    ___
#   / _ \  / _ \
#  / / \ \/ / \ \
# ( (   )  (   ) )
#  \ \_/ /\ \_/ /
#   \_  ____  _/
#    /_/    \_\
#     pretzel
#

Thread.abort_on_exception = true

mongbat_dir = File.dirname(File.expand_path($0))
$:.push mongbat_dir

require "lib/util.rb"
require "lib/predator4_module.rb"
require "lib/ipv4address.rb"
require "lib/ipv6address.rb"
require "externals/validator.rb"
require "externals/generator.rb"
require "thread.rb"
require "timeout.rb"
require "fileutils"
require 'digest/sha2'
require "cgi"
require "base64"

# the predator/evader binary should be in the same directory as this file
$predator4_binary = mongbat_dir + "/predator4"
if not File.exists?($predator4_binary)
  $predator4_binary = mongbat_dir + "/evader"
  if not File.exists?($predator4_binary)
    # ok, maybe its in the path
    if `which predator4`.strip.size > 0
      $predator4_binary = `which predator4`.strip
    else
      if `which evader`.strip.size > 0
        $predator4_binary = `which evader`.strip
      else
        log "Cannot find a predator4/evader binary (not in #{mongbat_dir} and not in path)", ERROR
        exit(-1)
      end
    end
  end
end

$version = nil
$licensed_to = nil
$attack_names = Predator4Module::Parser.get_attacks($predator4_binary)
`#{$predator4_binary} -v`.split(/\n/).each do |line|
  line.strip!
  if line =~ /^(Predator4|Evader) version (.+)$/
    $version = $2
  elsif line =~ /^Licensed to "(.+)"$/
    $licensed_to = $1
  end
end
if $licensed_to.nil?
  log "Using binary #{$predator4_binary} version #{$version}"
else
  log "Using binary #{$predator4_binary} version #{$version} licensed to #{$licensed_to}"
end
$predator4 = Predator4Module::Interface.new($predator4_binary)
$predator4_basename = File.basename($predator4_binary)

workers = 1
iface = nil
attacker = nil
victim = nil
gw = nil
time = 60
use_evasions = []
disable_evasions = []
mode = "random"
$default_attack = $attack_names.first
$show_remote = false
$attack_names.each do |name|
  if name =~ /^http_.+_server$/
    $show_remote = true
    break
  end
end
attack = $default_attack
$recdir = nil
$passthrough = nil
$min_evasions = 1
$max_evasions = -1
$stop_on_success = false
$start_index = 0
$end_index = -1
$dst_port = nil
$mask = nil
$remote_browser = "firefox"
$remote_method = "selenium"
$use_stages = true
$all_options = false
$log_only_failures = false
$log_everything = false
external_files = nil
external_generator = nil
$ips_per_worker = 1
$uid = nil
rand_seed = nil
$payload_obfuscation = true
min_evasions_modified = false
max_evasions_modified = false

def print_help()
  puts "#{File.basename($0)} - uses solo/dual/random evasions to attack target host"
  puts "Options:"
  puts "\t--mode=(solo|dual|triple|random)\t\tMode of attack, solo for each supported evasion with some options, dual for combinations of two and random (default) for random options"
  puts "\t--attack=(#{$attack_names.join("|")})\tAttack to use (default is #{$default_attack})"
  puts "\t--iface=<interface>\t\t\t\tInterface from which the attacks should originate"
  puts "\t--attacker=<src ip>(/mask)?\t\t\tStarting source IP for attackers; First worker will use this and if more are configured, they will use the following"
  puts "\t--victim=<dst ip>\t\t\t\tDestination IP for the attackers; Expecting correct host and vulnerable service on default port"
  if $show_remote
    puts "\t--remote_method=(local|selenium|drb)\t\tThe method to control the browsers (default is #{$remote_method})"
    puts "\t--remote_browser=(firefox|wget|ie)\t\tThe browser to use on the victim (default is #{$selenium_browser})"
  end
  #puts "\t--mask=<netmask or prefix>\t\t\tNetmask for IPv4 in CIDR notation, prefix length for IPv6"
  puts "\t--gw=<gw ip>\t\t\t\t\tGateway address if the victim is not in the local network (defaults to empty)"
  puts "\t--time=<time in seconds, default 60s>\t\tTime in seconds - stop attacking once time is up (--mode=random)"
  puts "\t--workers=<worker count, default 1>\t\tUse this many workers (and source IP addresses) to do the attacking"
  puts "\t--use_evasions=<evasion>(,evasion)*\t\tUse only these evasions"
  puts "\t--disable_evasions=<evasion>(,evasion)*\t\tDo not use these evasions"
  puts "\t--check_victim=(true|false)\t\t\tCheck that victim allows legal traffic without evasions before attacking (default true)"
  puts "\t--verify_delay=<delay>\t\t\t\tTime in ms to wait before attempting to verify shell port (default: #{$verify_delay}ms)"
  puts "\t--record=<recdir>\t\t\t\tRecord the attacks to dirname in pcap format"
  puts "\t--min_evasions=<min evasions>\t\t\tMinimum evasions for random mode (default: #{$min_evasions})"
  puts "\t--max_evasions=<max evasions>\t\t\tMaximum evasions for random mode (0 for unlimited) (default: #{$max_evasions})"
  puts "\t--index=<begin(-end)?>\t\t\t\tStart and optional end index for solo and dual mode"
  puts "\t--stop_on_success\t\t\t\tStop if an attack is successful"
  puts "\t--payload=<shell|..>\t\t\t\tPayloads types. Defaults to 'shell'. Some payload cannot be checked for success"
  puts "\t--disable_payload_obfuscation\t\t\tDisable payload obfuscation - use with caution"
  puts "\t--use_stages=<true|false>\t\t\tUse stages when available. Defaults to #{$use_stages}"
  puts "\t--all_options=<true|false>\t\t\tEnable use of all options (dangerous). Defaults to #{$all_options}"
  puts "\t--validator=<validator>(,validator)*\t\tUse this ruby code validator to evaluate whether the combination is valid (dangerous)"
  puts "\t--randseed=<randseed>\t\t\t\tThis sets the base64 randseed to allow for some repeatability"
  puts "\t--passthrough\t\t\t\t\tPass remaining unknown arguments directly to #{$predator4_basename}"
  puts "Example:"
  puts "#{$0} --attack=#{$default_attack} --iface=eth1 --attacker=10.0.0.100 --workers=16 --victim=10.0.0.3 --time=3600"
  exit
end

S_TRUE = "true"
S_FALSE = "false"
S_H = "-h"
S_HELP = "--help"
R_WORKERS = /^--workers=(\d+)$/
R_ATTACKER = /^--attacker=(\S+)$/
R_SRC_IP = /^--src_ip=(\S+)$/
R_VICTIM = /^--victim=(\S+)$/
R_DST_IP = /^--dst_ip=(\S+)$/
R_REMOTE_METHOD = /^--remote_method=(local|selenium|drb)$/
R_REMOTE_BROWSER = /^--remote_browser=(\S+)$/
R_GW = /^--gw=(\S+)$/
R_TIME = /^--time=(\d+)$/
R_IFACE = /^--iface=(\S+)$/
R_IF = /^--if=(\S+)$/
R_MODE = /^--mode=(solo|dual|triple|random)$/
R_ATTACK = /^--attack=(#{$attack_names.join("|")})$/
R_USE_EVASIONS = /^--use_evasions=(.+)$/
R_DISABLE_EVASIONS = /^--disable_evasions=(.+)$/
R_CHECK_VICTIM = /^--check_victim=(true|false)$/
R_RECORD = /^--record=(.+)$/
R_MIN_EVASIONS = /^--min_evasions=(\d+)$/
R_MAX_EVASIONS = /^--max_evasions=(\d+)$/
R_INDEX = /^--index=(\d+(-\d+)?)$/
R_STOP_ON_SUCCESS = /^--stop_on_success$/
R_PAYLOAD = /^--payload=(.+)$/
R_PASSTHROUGH = /^--passthrough$/
R_MASK = /^--mask=(\d+)$/
R_PREFIX = /^--prefix=(\d+)$/
R_VICTIM_PORT = /^--victim_port=(\d+)$/
R_DST_PORT = /^--dst_port=(\d+)$/
R_USE_STAGES = /^--use_stages=(true|false)$/
R_ALL_OPTIONS = /^--all_options=(true|false)$/
R_VALIDATOR = /^--validator=(.+)$/
R_GENERATOR = /^--generator=(.+)$/
R_IPS_PER_WORKER = /^--ips_per_worker=(\d+)$/
R_RANDSEED = /^--randseed=(.+)$/
R_LOG_ONLY_FAILURES = /^--log_only_failures$/
R_LOG_EVERYTHING = /^--log_everything$/
R_DISABLE_PAYLOAD_OBFUSCATION = /^--disable_payload_obfuscation$/
R_VERIFY_DELAY = /^--verify_delay=(\d+)$/
R_COMMA = /,/
R_ATTACKER_MASK = /^(.+)\/(\d+)$/
R_UID = /^--uid=(.+)$/

ARGV.each do |arg|
  if arg == S_H or arg == S_HELP
    print_help()
  elsif arg =~ R_WORKERS
    workers = $1.to_i
  elsif arg =~ R_ATTACKER or arg =~ R_SRC_IP
    attacker = $1
    if attacker =~ R_ATTACKER_MASK
      $mask = $2.to_i
      attacker = $1
    end
  elsif arg =~ R_VICTIM or arg =~ R_DST_IP
    victim = $1
  elsif arg =~ R_REMOTE_METHOD
    $remote_method = $1
  elsif arg =~ R_REMOTE_BROWSER
    $remote_browser = $1
  elsif arg =~ R_GW
    gw = $1
  elsif arg =~ R_TIME
    time = $1.to_i
  elsif arg =~ R_IFACE or arg =~ R_IF
    iface = $1
  elsif arg =~ R_MODE
    mode = $1
  elsif arg =~ R_ATTACK
    attack = $1
  elsif arg =~ R_USE_EVASIONS
    parts = $1.split(/,/)
    max = parts.size
    i = 0
    use_evasions = []
    while true do
      if (i+1) < max and parts[i] =~ /^\[.+/ and parts[i+1] =~ /.+\].+/
        ev = parts[i] + "," + parts[i+1]
        i += 1
      else
        ev = parts[i]
      end
      use_evasions.push ev
      i += 1
      break if i >= max
    end
  elsif arg =~ R_DISABLE_EVASIONS
    parts = $1.split(/,/)
    max = parts.size
    i = 0
    disable_evasions = []
    while true do
      if (i+1) < max and parts[i] =~ /^\[.+/ and parts[i+1] =~ /.+\].+/
        ev = parts[i] + "," + parts[i+1]
        i += 1
      else
        ev = parts[i]
      end
      disable_evasions.push ev
      i += 1
      break if i >= max
    end
  elsif arg =~ R_CHECK_VICTIM
    # PORT_USAGE
    # controls how many ports will be used at max per attack
    # 1 for the attack itself and PORT_USAGE-1 for clean runs
    # to validate that the windows machine is running
    # using only 1 for windows checking will destroy the throttiling
    # ability from overload; using many will result in running out
    # of ports quickly which in turn will lead to long periods
    # of waiting
    if $1 == S_FALSE
      log "Victim check disabled - will NOT notice if victim is no longer running"
      PORT_USAGE = 1
    else
      PORT_USAGE = 3
    end
  elsif arg =~ R_RECORD
    $recdir = File.expand_path($1)
  elsif arg =~ R_MIN_EVASIONS
    $min_evasions = $1.to_i
    min_evasions_modified = true
  elsif arg =~ R_MAX_EVASIONS
    $max_evasions = $1.to_i
    max_evasions_modified = true
  elsif arg =~ R_INDEX
    index = $1.split(/-/)
    $start_index = index[0].to_i
    if not index[1].nil?
      $end_index = index[1].to_i
    end
  elsif arg =~ R_STOP_ON_SUCCESS
    $stop_on_success = true
  elsif arg =~ R_PAYLOAD
    $payload = $1
  elsif arg =~ R_PASSTHROUGH
    $passthrough = []
  elsif arg =~ R_MASK or arg =~ R_PREFIX
    $mask = $1.to_i
  elsif arg =~ R_VICTIM_PORT or arg =~ R_DST_PORT
    $dst_port = $1.to_i
  elsif arg =~ R_USE_STAGES
    $use_stages = ($1==S_TRUE)
  elsif arg =~ R_ALL_OPTIONS
    $all_options = ($1==S_TRUE)
  elsif arg =~ R_VALIDATOR
    external_files = $1.split(R_COMMA)
  elsif arg =~ R_GENERATOR
    external_generator = $1
  elsif arg =~ R_IPS_PER_WORKER
    $ips_per_worker = $1.to_i
  elsif arg =~ R_RANDSEED
    rand_seed = $1
  elsif arg =~ R_LOG_ONLY_FAILURES
    $log_only_failures = true
  elsif arg =~ R_LOG_EVERYTHING
    $log_everything = true
  elsif arg =~ R_DISABLE_PAYLOAD_OBFUSCATION
    $payload_obfuscation = false
  elsif arg =~ R_VERIFY_DELAY
    $verify_delay = $1.to_i
  elsif arg =~ R_UID
    $uid = $1
  else
    if not $passthrough.nil?
      $passthrough.push arg
    else
      puts "Unhandled argument #{arg}"
      print_help()
    end
  end
end

# construct the attacker uid
$attacker_uid = "mongbat_#{Process.pid}"
if not $uid.nil?
  $attacker_uid += "_#{$uid}"
  log "mongbat uid specified: #{$uid} -> attacker uid is #{$attacker_uid}", DEBUG
  # check that no wild mongbats with our uid are running around
  `ps -aef | grep -v grep | grep ruby | grep mongbat | grep #{$uid}`.strip.split(/\n/).each do |line|
    pid = line.split(/\s+/, 3)[1]
    next if pid.to_i == Process.pid
    cmd = "kill -9 #{pid} 2>&1"
    log "Killing previous mongbat: #{cmd}: " + `#{cmd}`, DEBUG
  end
  # then check that no children of this uid are around
  `ps -aef | grep -v grep | grep #{File.basename($predator4_binary)} | grep #{$uid}`.strip.split(/\n/).each do |line|
    pid = line.split(/\s+/, 3)[1]
    cmd = "kill -9 #{pid} 2>&1"
    log "Killing previous mongbat children #{File.basename($predator4_binary)}: #{cmd}: " + `#{cmd}`, DEBUG
  end
end

if iface.nil?
  log "Please provide an interface with --iface", ERROR
  print_help()
end

if attacker.nil?
  log "Please provide an attacker IP address with --attacker", ERROR
  print_help()
end

if victim.nil?
  log "Please provide a victim IP address with --victim", ERROR
  print_help()
end


gw_ip = nil
if attacker =~ /:/
  attacker_ip = IPForge::IPv6Address.new(attacker)
  victim_ip = IPForge::IPv6Address.new(victim)
  if not gw.nil?
    gw_ip = IPForge::IPv6Address.new(gw)
  end
  ipv6 = true
else
  attacker_ip = IPForge::IPv4Address.new(attacker)
  victim_ip = IPForge::IPv4Address.new(victim)
  if not gw.nil?
    gw_ip = IPForge::IPv4Address.new(gw)
  end
  ipv6 = false
end

if not defined?(PORT_USAGE)
  PORT_USAGE = 3
end

$host_config = Predator4Module::HostConfig.new()
$host_config.parse()
host_interface = $host_config.interfaces[iface]
if not host_interface.nil?
  if not host_interface.carrier
    puts "Host interface #{iface} has no carrier"
    print_help()
  elsif not host_interface.up
    puts "Host interface #{iface} is not UP"
    print_help()
  else
    attacker_ips = []
    ip = attacker_ip
    $ips_per_worker.times do
      workers.times do
        attacker_ips.push ip
        ip += 1
      end
    end
    attacker_ips.each do |ip|
      if victim_ip == ip
        puts "victim IP #{ip} is among the attacker IP addresses (#{attacker_ips.join(", ")})"
        print_help()
      end
      if gw_ip == ip
        puts "Gateway IP #{ip} is among the attacker IP addresses (#{attacker_ips.join(", ")})"
        print_help
      end
      if ipv6
        host_interface.ipv6_addresses.each do |host_ip|
          if host_ip == ip
            puts "Local host IP #{ip} is among the attacker IP addresses (#{attacker_ips.join(", ")})"
            print_help()
          end
        end
      else
        host_interface.ipv4_addresses.each do |host_ip|
          if host_ip == ip
            puts "Local host IP #{ip} is among the attacker IP addresses (#{attacker_ips.join(", ")})"
            print_help()
          end
        end
      end
    end

    # do clean checks for all attacker addresses to see that routing works etc
    # instead of doing all, do first and last and maybe (later) some others as well
    first_ip = attacker_ips.first
    last_ip = attacker_ips.last
    [first_ip, last_ip].each do |ip|
      if attack =~ /_server/
        # TODO: consider what should happen here
      else
        nc = Predator4Module::NetworkConfig.new
        nc.iface = iface
        nc.src_ip = ip
        nc.dst_ip = victim_ip
        nc.gw_ip = gw
        nc.mask = $mask
        if not $dst_port.nil?
          nc.dst_port = $dst_port
        end
        log("Checking connection from #{ip} to #{victim_ip}", DEBUG)
        code, explanation, result, cmd, recname = $predator4.clean(attack, nc, 2)
        if code == 300 # TCP connection failed.
          puts "Could not connect (Host/port is down) using #{cmd}: #{explanation}."
          print_help()
        end
        log("Connection from #{ip} to #{victim_ip} OK", DEBUG)
      end
    end
  end
else
  puts "Host has no such interface #{iface} (available: #{$host_config.interfaces.keys.join(", ")})"
  print_help()
end

if mode =~ /^(solo|dual|triple)$/
  if min_evasions_modified
    puts "Cannot configure min_evasions for mode #{mode}"
    print_help()
  elsif max_evasions_modified
    puts "Cannot configure max_evasions for mode #{mode}"
    print_help()
  end
end

if attack =~ /_server$/
  if $remote_method == "selenium"
    require "selenium-webdriver"
  elsif $remote_method == "drb"
    require "drb/drb"
  end
end

if rand_seed.nil?
  r = srand(42)
  # truncate to 64bits
  rand_string = [r].pack("Q")
  rand_int = rand_string.unpack("Q")[0]
  srand(rand_int)
  rand_seed = Base64::encode64(rand_string).chomp
  log "Using rand seed #{Base64::encode64([srand(rand_int)].pack("Q"))}"
else
  rand_string = Base64::decode64(rand_seed)
  rand_int = rand_string.unpack("Q")[0]
  if rand_int.nil?
    log "Please provide a valid base64 string as rand_seed (and not #{rand_seed})"
    exit
  end
  srand(rand_int)
  log "Using rand seed #{Base64::encode64([srand(rand_int)].pack("Q"))}"
end

if not $recdir.nil?
  FileUtils.mkdir_p($recdir)
end

if $predator4.running?
  log "#{$predator4_basename} is already running ; this may cause VICTIM CHECK FAILED messages!", WARN
end

$external_classes = []
def register_external(external)
  $external_classes.push external
end

if external_files.nil?
  default_validator = "externals/#{attack}_validator.rb"
  if File.exists?(default_validator)
    log "loading default validator #{default_validator}", DEBUG
    require File.expand_path(default_validator)
  end
else
  external_files.each do |external_file|
    if not File.exists?(external_file)
      log "No such external validator #{external_file}", ERROR
      print_help()
    end
    log "loading #{external_file}", INFO
    require File.expand_path(external_file)
  end
end

$validators = []
$external_classes.each do |external_class|
  $validators.push (external_class.method(:new)).call()
end

$validators.each do |external|
  log "External Validator: #{external.to_s}", INFO
end
if $validators.size == 0
  $validators.push Validator.new
end

def valid_evasion?(evasion, option_hash)
  $validators.each do |validator|
    return false if not validator.valid_evasion?(evasion, option_hash)
  end
  return true
end

def valid_combo?(configured_evasion)
  if $min_evasions > configured_evasion.size
    return false
  end
  if $max_evasions > -1 and configured_evasion.size > $max_evasions
    return false
  end
  $validators.each do |validator|
    return false if not validator.valid_combo?(configured_evasion)
  end
  return true
end

def register_generator(gen)
  if not $external_generator_class.nil?
    raise "Double registration of external generator"
  end
  $external_generator_class = gen
end

if not external_generator.nil?
  if not File.exists?(external_generator)
    log "No such external generator #{external_generator}", ERROR
    print_help()
  end
  log "loading #{external_generator}", INFO
  require File.expand_path(external_generator)
end

class TestResult
  def initialize(mode, attack, time, workers, iface, attacker, victim, gw, use_evasion, disable_evasions, stop_on_success)
    @mode = mode
    @attack = attack
    @test_duration = time
    @workers = workers
    @iface = iface
    @attacker = attacker
    @victim = victim
    @gw = gw
    @use_evasion = use_evasion
    @disable_evasions = disable_evasions
    @stop_on_success = stop_on_success
    @lock = Mutex.new
    @total_attempts = 0
    @clean_attempts = 0
    @exploit_attempts = 0
    @clean_successes = 0
    @exploit_successes = []
    @start_time = Time.new
  end

  def inc_total_attempts
    @lock.synchronize do
      @total_attempts += 1
    end
  end
  def inc_clean_attempts
    @lock.synchronize do
      @clean_attempts += 1
    end
  end
  def inc_exploit_attempts
    @lock.synchronize do
      @exploit_attempts += 1
    end
  end
  def add_exploit_successes(evasion, randseed, recname)
    @lock.synchronize do
      @exploit_successes.push([evasion, randseed, recname])
    end
  end
  def inc_clean_successes
    @lock.synchronize do
      @clean_successes += 1
    end
  end

  def txt_table(head, result_table, head_rows = 1, colors = false, highlight = false, linkmap = {})
    result = head + "\n"
    lengths = []
    result_table.first.length.times do |i|
      lengths.push(0)
      result_table.each do |row|
        lengths[i] = row[i].to_s.length if row[i].to_s.length > lengths[i]
      end
    end
    result_table.each do |row|
      if row.length > result_table.first.length
        puts "row length differs from header length: " + row.join(", "), ALERT
      end
      row.length.times do |i|
        result += sprintf("%#{lengths[i]}s\t", row[i])
      end
      result += "\n"
    end
    return result
  end

  def html_table(head, result_table, head_rows = 1, colors = false, highlight = false, linkmap = {})
    result = "<h3>#{head}</h3>"
    result += '<table style="width: 100%;" border="1" cellpadding="1" cellspacing="0">'
    row_count = 0
    color_on = true
    result_table.each do |row|
      first = row_count < head_rows
      bg_color = first ? S_BG_COLOR_WHITE : (color_on ? S_BG_COLOR_LGRAY : S_BG_COLOR_GRAY)
      result.print bg_color
      if colors
        color_sort = row[1..-1].sort{|x, y| x.to_f <=> y.to_f}
        best = color_sort[-1]
        worst = color_sort[0]
      end
      row_first = true
      row.each do |element|
        result.print S_TD
        result.print S_B if first
        if not first and not row_first
          if colors
            if element == best
              result.print S_FONT_GREEN
            elsif element == worst
              result.print S_FONT_RED
            else
              result.print S_FONT_BLUE
            end
            result.print CGI::escapeHTML(element.to_s)
            result.print S_FONT_END
          elsif highlight
            if element =~ R_OK
              result.print S_FONT_GREEN
            elsif element =~ R_FAIL
              result.print S_FONT_RED
            else
              result.print S_FONT_BLUE
            end
            result.print CGI::escapeHTML(element.to_s)
            result.print S_FONT_END
          else
            result.print CGI::escapeHTML(element.to_s)
          end
        else
          if first and not row_first and linkmap.has_key?(element)
            result.print "<a href='#{linkmap[element]}'>" + CGI::escapeHTML(element.to_s) + '</a>'
          else
            result.print CGI::escapeHTML(element.to_s)
          end
        end
        result.print S_BE if first
        result.print S_TDE
        row_first = false
      end
      result.puts "</tr>"
      row_count += 1
      color_on = (not color_on) if not first
    end
    result.puts "</table>"
    return result
  end

  def to_s
    puts "Printing test result"
    @end_time = Time.new
    if @attacker =~ /:/
      @attacker_ip = IPForge::IPv6Address.new(@attacker)
      @victim_ip = IPForge::IPv6Address.new(@victim)
      @ipv6 = true
    else
      @attacker_ip = IPForge::IPv4Address.new(@attacker)
      @victim_ip = IPForge::IPv4Address.new(@victim)
      @ipv6 = false
    end
    result =
"Mongbat test report

Using #{$predator4_binary} version #{$version}"
if not $licensed_to.nil?
  result += " licensed to #{$licensed_to}"
end

result += "

Started : #{@start_time}
Finished: #{@end_time}
Attack  : #{@attack}

Network setup:
Attackers: #{@attacker_ip.to_s}-#{(@attacker_ip+@workers*$ips_per_worker).to_s}
Victim   : #{@victim}
"
    if not @gw.nil?
      result += "Gateway  : #{@gw}"
    end

    result += "\n"
    result_table =
[
["", "Total", "Clean", "Exploit"],
["Attempts", @total_attempts, @clean_attempts, @exploit_attempts],
["Failures", @total_attempts - @exploit_successes.size, @clean_attempts - @clean_successes, @exploit_attempts - @exploit_successes.size],
["Success", @exploit_successes.size, @clean_successes, @exploit_successes.size]
]

    result += txt_table("Number summary", result_table)
    result += "\n"
    if @exploit_successes.size > 0
      result += "Exploits succeeded with the following evasions:\n"
      @exploit_successes.each do |evasions, randseed, record|
        evasions.each do |evasion, options|
          result += " --evasion=#{evasion.name}"
          evasion.options.each do |option|
            result += ",\"#{options[option.name]}\""
          end
        end
        result += " --randseed=#{randseed}"
        if not record.nil?
          result += " (PCAP: #{record})"
        end
        result += "\n"
      end
    end
    return result
  end
end

$interrupt = false
$interrupt_count = 0

def interrupt()
  $interrupt = true
  if $interrupt_count == 0
    puts "Interrupt registered, soft shutdown"
  elsif $interrupt_count == 1
    puts "Second interrupt registered, killing #{$predator4_basename}s"
    psout = `ps -afe | grep "#{$predator4_basename} " | grep "#{$attacker_uid} " | grep -v 'sh -c' | grep -v grep | grep -v ruby`.strip
    psout.split(/\n/).each do |line|
      pid = line.split(/\s+/, 3)[1]
      cmd = "kill -9 #{pid}"
      puts "#{cmd}: " + `#{cmd}`.chomp
    end
  elsif $interrupt_count == 2
    puts "Third interrupt, exiting.."
    exit
  else
    puts "Interrupt #{$interrupt_count}: going down hard"
    exit!
  end
  $interrupt_count += 1
end

Kernel.trap("SIGINT") {
  interrupt()
}

Kernel.trap("SIGKILL") {
  interrupt()
}

class MongBat
  def initialize(mode, attack, time, workers, iface, attacker, victim, gw, use_evasions = [], disable_evasions = [])
    @mode = mode
    @attack_name = attack
    @iface = iface
    @time = time
    @worker_count = workers
    @attacker = attacker
    @victim = victim
    $test_result = TestResult.new(mode, attack, time, workers, iface, attacker, victim, gw, use_evasions, disable_evasions, $stop_on_success)
    if @attacker =~ /:/
      @attacker_ip = IPForge::IPv6Address.new(@attacker)
      @victim_ip = IPForge::IPv6Address.new(@victim)
    else
      @attacker_ip = IPForge::IPv4Address.new(@attacker)
      @victim_ip = IPForge::IPv4Address.new(@victim)
      @ipv6 = false
    end
    @gw = gw

    @attack = $predator4.get_options(@attack_name, $use_stages, $all_options)

    if $payload.nil?
      # default
    elsif not @attack.payloads.include?($payload)
      puts "ERROR: Payload #{$payload} is not supported by #{@attack.short_name}, supported types are #{@attack.payloads.join(", ")}"
      exit
    else
      $passthrough = [] if $passthrough.nil?
      $passthrough.push "--payload=" + $payload
    end

    if not $payload_obfuscation
      log "Warning: Payload obfuscation is disabled with '--disable_payload_obfuscation' - this may trigger generic shellcode detectors and other non-vulnerability based detections"
    end
    delete_evasion = []
    @attack.evasions.each do |evasion|
      if @ipv6 and evasion.base_name =~ /^ipv4_/
        delete_evasion.push evasion
      elsif not @ipv6 and evasion.base_name =~ /^ipv6_/
        delete_evasion.push evasion
      end
    end
    delete_evasion.each do |evasion|
      @attack.evasions.delete(evasion)
    end

    @disable_evasions = {}
    @disable_evasions_order = []
    disable_evasions.each do |evasion_name|
      found = false
      @attack.evasions.each do |evasion|
        if evasion_name == evasion.name
          @disable_evasions[evasion_name] = true
          @disable_evasions_order.push evasion_name
          found = true
        elsif $use_stages and evasion_name == evasion.base_name
          @disable_evasions[evasion.name] = true
          @disable_evasions_order.push evasion.name
          found = true
        end
      end
      if not found
        log "Could not find evasion #{evasion_name} to disable"
        exit
      else
        log "Disabling evasion #{evasion_name}"
      end
    end

    @use_evasions = {}
    @use_evasions_order = []
    use_evasions.each do |evasion_name|
      found = false
      @attack.evasions.each do |evasion|
        if evasion_name == evasion.name
          @use_evasions[evasion_name] = true
          @use_evasions_order.push evasion_name
          found = true
        elsif $use_stages and evasion_name == evasion.base_name
          @use_evasions[evasion.name] = true
          @use_evasions_order.push evasion.name
          found = true
        end
      end

      if not found
        log "Could not find evasion #{evasion_name} to use"
        exit
      else
        log "Using evasion #{evasion_name}"
      end
    end
  end

  def start_evasion_generator()
    @test_queue = []
    @evasion_generator_thread = Thread.new do
      evindex = 0
      while(@test_running and not $interrupt)
        @thread_lock.synchronize do
          while @test_queue.size < (@worker_count*10)
            evindex += 1
            if $end_index != -1 and evindex > $end_index
              next # index out of bounds
            end

            # we have to generate the $start_index < evindex ones since some of the generators will not otherwise increment their internal state
            configured_evasion = $generator.generate(evindex)
            randseed = $generator.randseed(evindex)

            if $start_index > evindex
              next
            end

            if configured_evasion.nil?
              # wait for the end
              break
            else
              @test_queue.push [evindex, configured_evasion, randseed]
            end
          end
          @thread_condition_variable.signal
        end
        sleep 0.1
      end
    end
  end

  def start_janitor
    # this is where all the custom monitoring stuff should be
    @janitor_thread = Thread.new do
      while(@test_running and not $interrupt)
        sleep 1
        @thread_condition_variable.signal
      end
    end
  end

  def get_attack_options()
    evasion = nil
    index = nil
    randseed = nil
    @thread_lock.synchronize do
      if $generator.pull_mode
        evasion = $generator.generate($index)
        randseed = $generator.randseed($index)
        index = $index
        $index += 1
      else
        @thread_condition_variable.wait(@thread_lock)
        if @test_queue.size > 0
          index, evasion, randseed = @test_queue.pop
        end
      end
    end
    return index, evasion, randseed
  end
  
  # blocks until ip is available
  def allocate_server_ip()
    found = nil
    while found.nil?
      @pool_lock.synchronize do
        @server_ip_pool.each do |ip, free|
          if free
            found = ip
            @server_ip_pool[ip] = false
            break
          end
        end
        if not found
          @pool_condition_variable.wait(@pool_lock)
        end
      end
    end
    return found
  end

  def release_server_ip(ip)
    @pool_lock.synchronize do
      @server_ip_pool[ip] = true
      @pool_condition_variable.signal()
    end
  end
  
  def start_workers
    $driver_by_ip = {}
    @worker_count.times do |i|

      attacker_base = @attacker_ip + i * $ips_per_worker
      attacker_shell_port = 10000 + i
      attacker_ip = attacker_base

      if @attack_name =~ /_server$/
        if $remote_method == "selenium"
          driver = Selenium::WebDriver.for( :remote, :desired_capabilities => $remote_browser, :url => "http://#{@victim}:4444/wd/hub" )
          $driver_by_ip[attacker_ip] = driver
        elsif $remote_method == "drb"
          drb_interface = DRbObject.new(nil, "druby://#{@victim}:7331")
          drb_interface.timeout = 8
          if $remote_browser == "wget"
            drb_interface.browser = "wget --tries=1 --quiet --timeout=8 --output-document=-"
          else
            drb_interface.browser = $remote_browser
          end
          $driver_by_ip[attacker_ip] = drb_interface
        end
      end
      t = Thread.new do
        nc = Predator4Module::NetworkConfig.new
        nc.iface = @iface
        nc.src_ip = attacker_ip
        nc.dst_ip = @victim
        nc.gw_ip = @gw
        nc.mask = $mask
        if not $dst_port.nil?
          nc.dst_port = $dst_port
        end

        while(@test_running and not $interrupt) do
          index, evasion, randseed = get_attack_options()
          if not evasion.nil?
            code = -1
            explanation = ""
            if @attack_name =~ /_server/
              code, explanation, result, cmd, recname = $predator4.server_attack(@attack_name, nc, evasion, $dst_port, $recdir, 10, $driver_by_ip[attacker_ip], randseed, $passthrough)
            else
              if @attack.extra_options_by_name[ "bindport" ].nil?
                code, explanation, result, cmd, recname = $predator4.attack(@attack_name, nc, evasion, nil, nil, PORT_USAGE, $recdir, 300, randseed, $passthrough)
              else
                code, explanation, result, cmd, recname = $predator4.attack(@attack_name, nc, evasion, nil, attacker_shell_port, PORT_USAGE, $recdir, 300, randseed, $passthrough)
              end
            end

            attacker_ip += 1
            if((attacker_ip - $ips_per_worker) == attacker_base)
              attacker_ip = attacker_base
            end
            @attack_count_lock.synchronize do
              @attack_count += 1
            end

            nc.src_ip = attacker_ip
            log_result(code, explanation, result, cmd, evasion, attacker_ip, recname)
            $generator.result(code, index)
          else
            sleep 0.1
          end
          @thread_condition_variable.signal
        end
        begin
          if not driver.nil?
            driver.quit
          end
        rescue Timeout::Error
        end
      end
      @worker_threads.push t
    end
  end

  def run()
    start_time = Time.new
    test_time = @time
    abort = false
    $map_lock = Mutex.new
    $map = File.new("mongbat_result_codes.txt", "w")

    if $external_generator_class.nil?
      if @mode == "random"
        require "externals/random_generator.rb"
        $generator = RandomGenerator.new(@attack, @use_evasions, @disable_evasions)
      else
        require "externals/combination_generator.rb"
        if @mode == "solo"
          $generator = CombinationGenerator.new(@attack, @use_evasions, @disable_evasions, 1)
        elsif @mode == "dual"
          $generator = CombinationGenerator.new(@attack, @use_evasions, @disable_evasions, 2)
        elsif @mode == "triple"
          $generator = CombinationGenerator.new(@attack, @use_evasions, @disable_evasions, 3)
        else
          raise "Unhandled mode #{@mode}"
        end
      end
    else
      $generator = ($external_generator_class.method(:new)).call(@attack, @use_evasions, @disable_evasions)
    end

    @total = $generator.count()

    start_threads()
   
    sleep 1
    if @total.nil?
      i = 0
      while(true)
        current_time = Time.new
        if((start_time+test_time) < current_time) or $interrupt
          break
        end
        i+=1
        if i.modulo(50) == 1
          log_no_header ""
          used_time = current_time - start_time
          attack_average = @attack_count / used_time
          log_nlf "#{@attack_count} runs averaging #{sprintf("%0.2f", attack_average)} runs / second ; progress: #{used_time.round}/#{test_time}"
        end
        if not @test_running
          log "Something wicked happened in one of the workers, aborting"
          abort = true
          break
        end
        sleep 0.1
      end
    else
      i = 0
      while(true)
        i += 1
        if $interrupt
          break
        end
        if i.modulo(50) == 1
          log_no_header ""
          log_nlf "#{Time.new} #{@attack_count}/#{@total}"
        end
        @thread_condition_variable.signal
        if not @test_running
          log "Something wicked happened in one of the workers, aborting"
          abort = true
          break
        end
        if $generator.pull_mode
          break if $generator.completed
        else
          if @test_queue.size == 0
            break
          end
        end
        sleep 0.1
      end
    end
    stop_threads()
    log_no_header ""
    log "Done."
    $map.close
    if abort
      return false
    end
    return true
  end

  def start_threads()
    @test_running = true
    @attack_count = $start_index
    @worker_threads = []
    @thread_lock = Mutex.new
    @thread_condition_variable = ConditionVariable.new
    @attack_count_lock = Mutex.new

    puts "Starting evasions generator: #{$generator.to_s}"
    if not $generator.pull_mode()
      start_evasion_generator()
    else
      $index = 0
    end

    start_janitor()
    start_workers()
  end

  def stop_threads
    @test_running = false
    @janitor_thread.join
    @worker_threads.each do |thread|
      thread.join
    end
  end

  def log_result(code, explanation, result, cmd, evasions, worker, recname)
    $test_result.inc_total_attempts
    $map_lock.synchronize do
      evasions.each do |ev|
        $map.print ev[1].values.join(",") + " "
      end
      $map.puts code.to_s
      $map.flush
    end
    if 200 <= code and code <= 400 # connect fail
      $test_result.inc_clean_attempts
      $test_result.inc_exploit_attempts
      $test_result.inc_clean_successes
    elsif code == 299 # client fail
      $test_result.inc_clean_attempts
    elsif code > 990 # clean fail
      $test_result.inc_clean_attempts
    else
      $test_result.inc_clean_attempts
      $test_result.inc_exploit_attempts
      $test_result.inc_clean_successes
    end

    if $log_everything
      str = "\n#{explanation} (#{worker}):\n#{cmd}\n#{result}"
      if not recname.nil?
        str += "(PCAP: #{recname})"
      end
      log str
    elsif $log_only_failures
      if not code == 0
        str = "\n#{explanation} (#{worker}):\n#{cmd}\n#{result}"
        if not recname.nil?
          str += "(PCAP: #{recname})"
        end
        log str
      else
        log_nlf "s"
      end
    else # normal logging: success and errors
      if 200 <= code and code <= 400 # connect fail
        log_nlf "."
      elsif code == 299 # client fail
        log_nlf "*"
      elsif code > 990 # clean fail
        log_nlf "C"
      else
        str = "\n#{explanation} (#{worker}):\n#{cmd}\n#{result}"
        if not recname.nil?
          str += "(PCAP: #{recname})"
        end
        log str
      end
    end

    if code == 0
      if result =~ /Using random seed (\S+)/
        randseed = $1
      end
      $test_result.add_exploit_successes(evasions, randseed, recname)

      if $stop_on_success
        log $test_result.to_s
        $stdout.flush
        $stderr.flush
        exit! # not exactly elegant eh?
      end
    end
  end

end

mongbat = MongBat.new(mode, attack, time, workers, iface, attacker, victim, gw, use_evasions, disable_evasions)
mongbat.run()
log $test_result.to_s
