require File.expand_path(File.join(File.dirname(__FILE__), 'nexpose.rb'))
require File.expand_path(File.join(File.dirname(__FILE__), 'util.rb'))
require 'rexml/document'
require 'hirb'
require 'ostruct'
require 'optparse'
=begin
require 'rubygems'
require 'nexpose'
=end


class Options
    def self.parse(args)
      options = OpenStruct.new

      option_parser = OptionParser.new  do |option_parser|
        option_parser.on("-h HOST", "The network address of the NeXpose instance - Required") {|arg| options.host=arg}
        option_parser.on("-u USER", "The user name - Required") {|arg| options.user=arg}
        option_parser.on("-p PASSWORD", "The password - Required") {|arg| options.password=arg}
        option_parser.on("--n1", "List all the asset groups") {|arg| options.key=1}
        option_parser.on("--n2 GI", "List an asset groups config. The GI (group id or group name is required)") do |arg|
          options.key=2
          options.args=arg
        end
        option_parser.on("--n3 GID", "Start an asset group scan. The GID (group id is required)") do |arg|
          options.key=3
          options.args=arg
        end
        option_parser.on_tail("--help", "Help") do
          puts option_parser
          exit 0
        end
      end

      begin
       option_parser.parse!(args)
        rescue OptionParser::ParseError => e
            puts "#{e}\n\n#{option_parser}"
            exit(1)
      end

      options
    end
end

class APIUtil
  attr_accessor :user_name, :password, :host
  @api_connection = nil
  @@connected = false

  def self.is_connected
    @@connected
  end

  def initialize(user_name, password, host)
    @user_name = user_name
    @password = password
    @host = host
  end

  private
    def get_asset_group_id
      group_id = -1
      begin
        print("Enter the asset group ID:")
        group_id = gets().chomp()
        if (!Util.is_number?(group_id))
          raise Exception.new("Input is not a number")
        end
        accepted_input = true
      rescue Exception => e
        puts(e.message)
      end until accepted_input == true

      group_id;
    end

    def get_asset_scan_input()
      max_scans = -1
      scan_template_id = -1
      accepted_input = false
      group_id = get_asset_group_id

      accepted_input = false
      begin
        print("Enter the max amount of devices to scan (-1 for unlimitted):")
        max_scans = gets().chomp
        if (!Util.is_number?(max_scans))
          raise Exception.new("Input is not a number")
        end
        accepted_input = true
      rescue Exception => e
        puts(e.message)
      end until accepted_input == true

=begin
      accepted_input = false
      begin
        print("Enter the scan template to be used (-1 for default, l to list scan template IDs):")
        input = gets()
        scan_template_id = input.to_i
        accepted_input = true
      rescue Exception => e
        puts(e.message)
      end until accepted_input == true
=end

      {:group_id=>group_id, :max_scans=>max_scans}
    end

  public
    #-------------------------------------------------------------------
    # Logs in to NeXpose and sets a session key on the connector object.
    #-------------------------------------------------------------------
    def do_login()
      if (!@@connected)
        begin
          @api_connector = Nexpose::Connection.new(host, user_name, password)
          if (@api_connector.login)
            @@connected = true
          end
        rescue Exception => e
          puts(e.message)
        end
      end
    end

    #------------------------------------------------------
    # Prints all asset group information in tabular format:
    # | site_id | device_id | address | riskfactor |
    #------------------------------------------------------
    def print_asset_group_info(group_id)
      group_id = group_id != -1 ? group_id : get_asset_group_id
      group_configs = @api_connector.asset_group_config(group_id)
      puts("\nASSET GROUP INFO (id: #{group_id})")
      puts Hirb::Helpers::AutoTable.render(group_configs, :fields => [:site_id, :device_id, :address, :riskfactor])
    end

    #----------------------------------------------------------------
    # Prints asset group configuration information in tabular format:
    # | site_id | device_id | address | riskfactor |
    #----------------------------------------------------------------
    def print_asset_groups()
      res = @api_connector.asset_groups_listing
      puts("\nASSET GROUPS:")
      puts Hirb::Helpers::AutoTable.render(res, :fields => [:asset_group_id, :name, :description, :risk_score])
    end

    #----------------------------------------------------------------
    # Starts an asset group scan for a particular group ID
    #
    # group_input: A Hash containing the group_id ie: {group_id => 1}
    #----------------------------------------------------------------
    def start_asset_group_scan(group_input)
      input = group_input != -1 ? group_input : get_asset_scan_input
      if (group_input != -1 and (not Util.is_number? input[:group_id]))
        group_name = input[:group_id].chomp
        res = @api_connector.asset_groups_listing
        found = false
        for asset_info in res
          if asset_info[:name].chomp.eql?(group_name)
            input[:group_id] = asset_info[:asset_group_id]
            found = true
            break
          end
        end

        if not found
          raise "Unable to find asset group with name: #{group_name}"
        end

      end

      group_configs = @api_connector.asset_group_config(input[:group_id])

      # First build a reverse map to ensure we will not be scanning a device more than once.
      device_to_site = Hash.new("")
      for group_config in group_configs
        key = group_config[:address]
        value = [group_config[:site_id], group_config[:device_id]].flatten
        device_to_site[key] = value;
      end

      # Combobulate to a map of site_id->addresses
      site_to_devices = Hash.new("")
      for address in device_to_site.keys
        key = device_to_site[address][0]
        value = device_to_site[address][1]
        if (site_to_devices.has_key?(key))
          site_to_devices[key] = [site_to_devices[key], value].flatten()
        else
          site_to_devices[key] = [value]
        end
      end

      # TODO: More intelligent logic needed, ie: scan already started
      for site in site_to_devices.keys
        res = @api_connector.site_device_scan_start(site, site_to_devices[site], nil)
        if (res)
          puts("Scan started scan ID: #{res[:scan_id]}, on engine ID: #{res[:engine_id]}")
        else
          puts("Scan failed for site #{site}")
        end
      end
    end

    def do_logout
      @api_connector.logout
    end

end


#---------------------
# MAIN EXECUTION BLOCK
#---------------------

begin

if ARGV.length > 0
  begin
    options = Options.parse ARGV
    api_util = ::APIUtil.new options.user, options.password, options.host
    api_util.do_login
    case options.key
      when 1 then api_util.print_asset_groups
      when 2 then api_util.print_asset_group_info options.args
      when 3 then api_util.start_asset_group_scan({:group_id => options.args, :max_scans => -1})
      else puts "Invalid Input!"
    end
  rescue Exception => e
    puts e
  end
  exit 0
else
  # Show Main Commands
  commands =
<<CMDS
  SELECT A COMMAND:
  1 - Print asset groups info
  2 - Print an asset group info
  3 - Scan an asset group
  q - Quit
CMDS

  while (!APIUtil.is_connected) do
    print("Enter the NeXpose host address:")
    address = gets().chomp()
    print("Enter the user name for this instance:")
    user_name = gets().chomp()
    print("Enter the password for this instance:")
    password = gets().chomp()
    api_util = ::APIUtil.new user_name, password, address
    api_util.do_login
  end

  input = ''
  while (input[0] != 'q') do
    puts("\n\n#{commands}")
    print("\nSelect Option:")
    input = gets()
    puts()

    if (!APIUtil.is_connected)
      api_util.do_login
    end

    case ( input[0] )
      when '1' then api_util.print_asset_groups
      when '2' then api_util.print_asset_group_info -1
      when '3' then api_util.start_asset_group_scan -1
      else if input[0] != 'q' then puts("Invalid Input!") end
    end
  end

  puts("Bye")
  if (APIUtil.is_connected)
    api_util.do_logout
  end

end

end