msfdb

#!/usr/bin/env ruby

require 'fileutils'
require 'io/console'
require 'json'
require 'net/http'
require 'net/https'
require 'open3'
require 'optparse'
require 'rex/socket'
require 'rex/text'
require 'securerandom'
require 'uri'
require 'yaml'
require 'pg'


include Rex::Text::Color

msfbase = __FILE__
while File.symlink?(msfbase)
  msfbase = File.expand_path(File.readlink(msfbase), File.dirname(msfbase))
end

$:.unshift(File.expand_path(File.join(File.dirname(msfbase), 'lib')))
$:.unshift(ENV['MSF_LOCAL_LIB']) if ENV['MSF_LOCAL_LIB']

require 'msfdb_helpers/pg_ctlcluster'
require 'msfdb_helpers/pg_ctl'
require 'msfdb_helpers/standalone'

require 'msfenv'

@script_name = File.basename(__FILE__)
@framework = File.expand_path(File.dirname(__FILE__))

@localconf = Msf::Config.get_config_root
@db = "#{@localconf}/db"
@db_conf = "#{@localconf}/database.yml"
@pg_cluster_conf_root = "#{@localconf}/.local/etc/postgresql"
@db_driver = nil

@ws_tag = 'msf-ws'
@ws_conf = File.join(@framework, "#{@ws_tag}.ru")
@ws_ssl_key_default = "#{@localconf}/#{@ws_tag}-key.pem"
@ws_ssl_cert_default = "#{@localconf}/#{@ws_tag}-cert.pem"
@ws_log = "#{@localconf}/logs/#{@ws_tag}.log"
@ws_pid = "#{@localconf}/#{@ws_tag}.pid"

@current_user = ENV['LOGNAME'] || ENV['USERNAME'] || ENV['USER']
@msf_ws_user = (@current_user || "msfadmin").to_s.strip
@ws_generated_ssl = false
@ws_api_token = nil

@components = %w(database webservice)
@environments = %w(production development)

@options = {
    # When the component value is nil, the user has not yet specified a specific component
    # It will later be defaulted to a more sane value
    component: nil,
    debug: false,
    msf_db_name: 'msf',
    msf_db_user: 'msf',
    msftest_db_name: 'msftest',
    msftest_db_user: 'msftest',
    db_host: '127.0.0.1',
    db_port: 5433,
    db_pool: 200,
    address: 'localhost',
    port: 5443,
    daemon: true,
    ssl: true,
    ssl_cert: @ws_ssl_cert_default,
    ssl_key: @ws_ssl_key_default,
    ssl_disable_verify: true,
    ws_env: ENV['RACK_ENV'] || 'production',
    retry_max: 10,
    retry_delay: 5.0,
    ws_user: nil,
    add_data_service: true,
    data_service_name: nil,
    use_defaults: false,
    delete_existing_data: true
}

def supports_color?
  return true if Rex::Compat.is_windows
  term = Rex::Compat.getenv('TERM')
  term and term.match(/(?:vt10[03]|xterm(?:-color)?|linux|screen|rxvt)/i) != nil
end

class String
  def bold
    substitute_colors("%bld#{self}%clr")
  end

  def underline
    substitute_colors("%und#{self}%clr")
  end

  def red
    substitute_colors("%red#{self}%clr")
  end

  def green
    substitute_colors("%grn#{self}%clr")
  end

  def blue
    substitute_colors("%blu#{self}%clr")
  end

  def cyan
    substitute_colors("%cya#{self}%clr")
  end

end

def pw_gen
  SecureRandom.base64(32)
end

def tail(file)
  begin
    File.readlines(file).last.to_s.strip
  rescue
    nil
  end
end

def status_db
  update_db_port

  case @db_driver.status
  when DatabaseStatus::RUNNING
    puts "Database started"
  when DatabaseStatus::INACTIVE
    puts "Database found, but is not running"
  when DatabaseStatus::NEEDS_INIT
    puts "Database found, but needs initialized"
  when DatabaseStatus::NOT_FOUND
    puts "No database found"
  end
end

def start_db
  case @db_driver.status
  when DatabaseStatus::NOT_FOUND
    print_error 'No database found.'
    return
  when DatabaseStatus::NEEDS_INIT
    print_error 'Has the database been initialized with "msfdb init" or "msfdb init --component database"?'
    return
  end

  update_db_port
  db_started = @db_driver.start

  if !db_started
    last_log = tail("#{@db}/log")
    puts last_log
    if last_log =~ /not compatible/
      puts 'Please attempt to upgrade the database manually using pg_upgrade.'
    end
    print_error 'Your database may be corrupt. Try reinitializing.'
  end
end

def stop_db
  update_db_port
  @db_driver.stop
end

def restart_db
  @db_driver.restart
end

def init_db
  case @db_driver.status
  when DatabaseStatus::RUNNING
    puts 'Existing database running'
    return
  when DatabaseStatus::INACTIVE
    puts 'Existing database found, attempting to start it'
    @db_driver.start
    return
  end

  if @db_driver.exists? && !@options[:delete_existing_data]
    if !load_db_config
      puts 'Failed to load existing database config. Please reinit and overwrite the file.'
      return
    end
  end

  # Generate new database passwords if not already assigned
  @msf_pass ||= pw_gen
  @msftest_pass ||= pw_gen

  @db_driver.init(@msf_pass, @msftest_pass)
  write_db_config

  puts 'Creating initial database schema'
  Dir.chdir(@framework) do
    @db_driver.run_cmd('bundle exec rake db:migrate')
  end
end

def load_db_config
  if File.file?(@db_conf)
    config = YAML.load(File.read(@db_conf))

    production = config['production']
    if production.nil?
      puts "No production section found in database config #{@db_conf}."
      return false
    end

    test = config['test']
    if test.nil?
      puts "No test section found in database config #{@db_conf}."
      return false
    end

    # get values for development and production
    @options[:msf_db_name] = production['database']
    @options[:msf_db_user] = production['username']
    @msf_pass = production['password']
    @options[:db_port] = production['port']
    @options[:db_pool] = production['pool']

    # get values for test
    @options[:msftest_db_name] = test['database']
    @options[:msftest_db_user] = test['username']
    @msftest_pass = test['password']
    return true
  end

  return false
end

def write_db_config
  # Write a default database config file
  Dir.mkdir(@localconf) unless File.directory?(@localconf)
  File.open(@db_conf, 'w') do |f|
    f.puts <<~EOF
      development: &pgsql
        adapter: postgresql
        database: #{@options[:msf_db_name]}
        username: #{@options[:msf_db_user]}
        password: #{@msf_pass}
        host: #{@options[:db_host]}
        port: #{@options[:db_port]}
        pool: #{@options[:db_pool]}

      production: &production
        <<: *pgsql

      test:
        <<: *pgsql
        database: #{@options[:msftest_db_name]}
        username: #{@options[:msftest_db_user]}
        password: #{@msftest_pass}
    EOF
  end

  File.chmod(0640, @db_conf)
end

def update_db_port
  if File.file?(@db_conf)
    config = YAML.load(File.read(@db_conf))
    if config["production"] && config["production"]["port"]
      port = config["production"]["port"]
      if port != @options[:db_port]
        puts "Using database port #{port} found in #{@db_conf}"
        @options[:db_port] = port
      end
    end
  end
end

def ask_yn(question, default: nil)
  loop do
    print "#{'[?]'.blue.bold} #{question} [#{default}]: "
    input = STDIN.gets.strip
    input = input.empty? ? default : input
    case input
    when /^[Yy]/
      return true
    when /^[Nn]/
      return false
    else
      puts 'Please answer yes or no.'
    end
  end
end

def ask_value(question, default)
  return default if @options[:use_defaults]

  print "#{'[?]'.blue.bold} #{question} [#{default}]: "
  input = STDIN.gets.strip
  if input.nil? || input.empty?
    return default
  else
    return input
  end
end

def ask_password(question)
  print "#{'[?]'.blue.bold} #{question}: "
  input = STDIN.noecho(&:gets).chomp
  print "\n"
  if input.nil? || input.empty?
    return pw_gen
  else
    return input
  end
end

def print_error(error)
  puts "#{'[!]'.red.bold} #{error}"
end

def delete_db
  stop_web_service
  @db_driver.delete
end

def reinit_db
  delete_db
  init_db
end

class WebServicePIDStatus
  RUNNING = 0
  INACTIVE = 1
  NO_PID_FILE = 2
end

class DatabaseStatus
  RUNNING = 0
  INACTIVE = 1
  NOT_FOUND = 2
  NEEDS_INIT = 3
end

def web_service_pid
  File.file?(@ws_pid) ? tail(@ws_pid) : nil
end

def web_service_pid_status
  if File.file?(@ws_pid)
    ws_pid = tail(@ws_pid)
    if ws_pid.nil? || !process_active?(ws_pid.to_i)
      WebServicePIDStatus::INACTIVE
    else
      WebServicePIDStatus::RUNNING
    end
  else
    WebServicePIDStatus::NO_PID_FILE
  end
end

def status_web_service
  ws_pid = web_service_pid
  status = web_service_pid_status
  if status == WebServicePIDStatus::RUNNING
    puts "MSF web service is running as PID #{ws_pid}"
  elsif status == WebServicePIDStatus::INACTIVE
    puts "MSF web service is not running: PID file found at #{@ws_pid}, but no active process running as PID #{ws_pid}"
  elsif status == WebServicePIDStatus::NO_PID_FILE
    puts "MSF web service is not running: no PID file found at #{@ws_pid}"
  end
end

def init_web_service
  if web_service_pid_status == WebServicePIDStatus::RUNNING
    puts "MSF web service is already running as PID #{web_service_pid}"
    return false
  end

  unless @options[:use_defaults]
    if @options[:ws_user].nil?
      @msf_ws_user = ask_value('Initial MSF web service account username?', @msf_ws_user)
    else
      @msf_ws_user = @options[:ws_user]
    end
  end

  if @options[:use_defaults]
    @msf_ws_pass = pw_gen
  elsif @options[:ws_pass].nil?
    @msf_ws_pass = ask_password('Initial MSF web service account password? (Leave blank for random password)')
  else
    @msf_ws_pass = @options[:ws_pass]
  end

  if should_generate_web_service_ssl && @options[:delete_existing_data]
    generate_web_service_ssl(key: @options[:ssl_key], cert: @options[:ssl_cert])
  end

  if start_web_service(expect_auth: false)
    if add_web_service_workspace && add_web_service_user
      output_web_service_information
    else
      puts 'Failed to complete MSF web service configuration, please reinitialize.'
      stop_web_service
    end
  end
end

def start_web_service_daemon(expect_auth:)
  if @db_driver.run_cmd("#{thin_cmd} start") == 0
    # wait until web service is online
    retry_count = 0
    response_data = web_service_online_check(expect_auth: expect_auth)
    is_online = response_data[:state] != :offline
    while !is_online && retry_count < @options[:retry_max]
      retry_count += 1
      if @options[:debug]
        puts "MSF web service doesn't appear to be online. Sleeping #{@options[:retry_delay]}s until check #{retry_count}/#{@options[:retry_max]}"
      end
      sleep(@options[:retry_delay])
      response_data = web_service_online_check(expect_auth: expect_auth)
      is_online = response_data[:state] != :offline
    end

    if response_data[:state] == :online
      puts "#{'success'.green.bold}"
      puts 'MSF web service started and online'
      return true
    elsif response_data[:state] == :error
      puts "#{'failed'.red.bold}"
      print_error 'MSF web service failed and returned the following message:'
      puts "#{response_data[:message].nil? || response_data[:message].empty? ? "No message returned." : response_data[:message]}"
    elsif response_data[:state] == :offline
      puts "#{'failed'.red.bold}"
      print_error 'A connection with the web service was refused.'
    end

    puts "Please see #{@ws_log} for additional webservice details."
    return false
  else
    puts "#{'failed'.red.bold}"
    puts 'Failed to start MSF web service'
    return false
  end
end

def start_web_service(expect_auth: true)
  unless File.file?(@ws_conf)
    puts "No MSF web service configuration found at #{@ws_conf}, not starting"
    return false
  end

  # check if MSF web service is already started
  ws_pid = web_service_pid
  status = web_service_pid_status
  if status == WebServicePIDStatus::RUNNING
    puts "MSF web service is already running as PID #{ws_pid}"
    return false
  elsif status == WebServicePIDStatus::INACTIVE
    puts "MSF web service PID file found, but no active process running as PID #{ws_pid}"
    puts "Deleting MSF web service PID file #{@ws_pid}"
    File.delete(@ws_pid)
  end

  print 'Attempting to start MSF web service...'

  unless File.file?(@options[:ssl_key])
    puts "#{'failed'.red.bold}"
    print_error "The SSL Key needed for the webservice to connect to the database could not be found at #{@options[:ssl_key]}."
    print_error 'Has the webservice been initialized with "msfdb init"  or "msfdb init --component webservice"?'
    return false
  end

  if @options[:daemon]
    start_web_service_daemon(expect_auth: expect_auth)
  else
    puts thin_cmd
    system "#{thin_cmd} start"
  end
end

def stop_web_service
  ws_pid = web_service_pid
  status = web_service_pid_status
  if status == WebServicePIDStatus::RUNNING
    puts "Stopping MSF web service PID #{ws_pid}"
    @db_driver.run_cmd("#{thin_cmd} stop")
  else
    puts 'MSF web service is no longer running'
    if status == WebServicePIDStatus::INACTIVE
      puts "Deleting MSF web service PID file #{@ws_pid}"
      File.delete(@ws_pid)
    end
  end
end

def restart_web_service
  stop_web_service
  start_web_service
end

def delete_web_service
  stop_web_service

  File.delete(@ws_pid) if web_service_pid_status == WebServicePIDStatus::INACTIVE
  if @options[:delete_existing_data]
    File.delete(@options[:ssl_key]) if File.file?(@options[:ssl_key])
    File.delete(@options[:ssl_cert]) if File.file?(@options[:ssl_cert])
  end
end

def reinit_web_service
  delete_web_service
  init_web_service
end

def generate_web_service_ssl(key:, cert:)
  @ws_generated_ssl = true
  if (File.file?(key) || File.file?(cert)) && !@options[:delete_existing_data]
    return
  end

  puts 'Generating SSL key and certificate for MSF web service'
  @ssl_key, @ssl_cert, @ssl_extra_chain_cert = Rex::Socket::Ssl.ssl_generate_certificate

  # write PEM format key and certificate
  mode = 'wb'
  mode_int = 0600
  File.open(key, mode) { |f| f.write(@ssl_key.to_pem) }
  File.chmod(mode_int, key)

  File.open(cert, mode) { |f| f.write(@ssl_cert.to_pem) }
  File.chmod(mode_int, cert)
end

def web_service_online_check(expect_auth:)
  msf_version_uri = get_web_service_uri(path: '/api/v1/msf/version')
  response_data = http_request(uri: msf_version_uri, method: :get,
                          skip_verify: skip_ssl_verify?, cert: get_ssl_cert)

  if !response_data[:exception].nil? && response_data[:exception].is_a?(Errno::ECONNREFUSED)
    response_data[:state] = :offline
  elsif !response_data[:exception].nil? && response_data[:exception].is_a?(OpenSSL::OpenSSLError)
    response_data[:state] = :error
    response_data[:message] = 'Detected an SSL issue. Please set the same options used to initialize the web service or reinitialize.'
  elsif !response_data[:response].nil? && response_data[:response].dig(:error, :code) == 401
    if expect_auth
      response_data[:state] = :online
    else
      response_data[:state] = :error
      response_data[:message] = 'MSF web service expects authentication. If you wish to reinitialize the web service account you will need to reinitialize the database.'
    end
  elsif !response_data[:response].nil? && !response_data[:response].dig(:data, :metasploit_version).nil?
    response_data[:state] = :online
  else
    response_data[:state] = :error
  end

  puts "web_service_online: expect_auth=#{expect_auth}, response_msg=#{response_data}" if @options[:debug]
  response_data
end

def add_web_service_workspace(name: 'default')
  # Send request to create new workspace
  workspace_data = { name: name }
  workspaces_uri = get_web_service_uri(path: '/api/v1/workspaces')
  response_data = http_request(uri: workspaces_uri, data: workspace_data, method: :post,
                               skip_verify: skip_ssl_verify?, cert: get_ssl_cert)
  response = response_data[:response]
  puts "add_web_service_workspace: add workspace response=#{response}" if @options[:debug]
  if response.nil? || response.dig(:data, :name) != name
    print_error "Error creating MSF web service workspace '#{name}'"
    return false
  end
  return true
end

def add_web_service_user
  puts "Creating MSF web service user #{@msf_ws_user}"

  # Generate new web service user password
  cred_data = { username: @msf_ws_user, password: @msf_ws_pass }

  # Send request to create new admin user
  user_data = cred_data.merge({ admin: true })
  user_uri = get_web_service_uri(path: '/api/v1/users')
  response_data = http_request(uri: user_uri, data: user_data, method: :post,
                          skip_verify: skip_ssl_verify?, cert: get_ssl_cert)
  response = response_data[:response]
  puts "add_web_service_user: create user response=#{response}" if @options[:debug]
  if response.nil? || response.dig(:data, :username) != @msf_ws_user
    print_error "Error creating MSF web service user #{@msf_ws_user}"
    return false
  end

  puts "\n#{'    ############################################################'.cyan}"
  print "#{'    ##              '.cyan}"
  print"#{'MSF Web Service Credentials'.cyan.bold.underline}"
  puts"#{'               ##'.cyan}"
  puts "#{'    ##                                                        ##'.cyan}"
  puts "#{'    ##        Please store these credentials securely.        ##'.cyan}"
  puts "#{'    ##    You will need them to connect to the webservice.    ##'.cyan}"
  puts "#{'    ############################################################'.cyan}"

  puts "\n#{'MSF web service username'.cyan.bold}: #{@msf_ws_user}"
  puts "#{'MSF web service password'.cyan.bold}: #{@msf_ws_pass}"

  # Send request to create new API token for the user
  generate_token_uri = get_web_service_uri(path: '/api/v1/auth/generate-token')
  response_data = http_request(uri: generate_token_uri, data: cred_data, method: :post,
                          skip_verify: skip_ssl_verify?, cert: get_ssl_cert)
  response = response_data[:response]
  puts "add_web_service_user: generate token response=#{response}" if @options[:debug]
  if response.nil? || (@ws_api_token = response.dig(:data, :token)).nil?
    print_error "Error creating MSF web service user API token"
    return false
  end
  puts "#{'MSF web service user API token'.cyan.bold}: #{@ws_api_token}"
  return true
end

def output_web_service_information
  puts "\n\n"
  puts 'MSF web service configuration complete'
  if @options[:add_data_service]
    data_service_name = @options[:data_service_name] || "local-#{@options[:ssl] ? 'https' : 'http'}-data-service"
    puts "The web service has been configured as your default data service in msfconsole with the name \"#{data_service_name}\""
  else
    puts "No data service has been configured in msfconsole."
  end
  puts ''
  puts 'If needed, manually reconnect to the data service in msfconsole using the command:'
  puts "#{get_db_connect_command}"
  puts ''
  puts 'The username and password are credentials for the API account:'
  puts "#{get_web_service_uri(path: '/api/v1/auth/account')}"
  puts ''

  if @options[:add_data_service]
    persist_data_service
  end
end

def run_msfconsole_command(cmd)
  # Attempts to run a the metasploit command first with the default env settings, and once again with the path set
  # to the current directory. This ensures that it works in an environment such as bundler
  if @db_driver.run_cmd(cmd) != 0
    # attempt to execute msfconsole in the current working directory
    if @db_driver.run_cmd(cmd, env: {'PATH' => ".:#{ENV["PATH"]}"}) != 0
      puts 'Failed to run msfconsole'
    end
  end
end

def persist_data_service
  puts 'Persisting http web data service credentials in msfconsole'
  # execute msfconsole commands to add and persist the data service connection
  cmd = "./msfconsole -qx '#{get_db_connect_command}; db_save; exit'"
  run_msfconsole_command(cmd)
end

def clear_default_data_service
  puts 'Clearing http web data service credentials in msfconsole'
  # execute msfconsole commands to clear the default data service connection
  cmd = "./msfconsole -qx 'db_disconnect --clear; exit'"
  run_msfconsole_command(cmd)
end

def get_db_connect_command
  data_service_name = "local-#{@options[:ssl] ? 'https' : 'http'}-data-service"
  if !@options[:data_service_name].nil?
    data_service_name = @options[:data_service_name]
  end

  # build db_remove and db_connect command based on install options
  connect_cmd = "db_connect"
  connect_cmd << " --name #{data_service_name}"
  connect_cmd << " --token #{@ws_api_token}"
  connect_cmd << " --cert #{@options[:ssl_cert]}" if @options[:ssl]
  connect_cmd << " --skip-verify" if skip_ssl_verify?
  connect_cmd << " #{get_web_service_uri}"
  connect_cmd
end

def get_web_service_uri(path: nil)
  uri_class = @options[:ssl] ? URI::HTTPS : URI::HTTP
  uri_class.build({host: get_web_service_host, port: @options[:port], path: path})
end

def get_web_service_host
  # user specified any address INADDR_ANY (0.0.0.0), return a routable address
  @options[:address] == '0.0.0.0' ? 'localhost' : @options[:address]
end

def skip_ssl_verify?
  @ws_generated_ssl || @options[:ssl_disable_verify]
end

def get_ssl_cert
  @options[:ssl] ? @options[:ssl_cert] : nil
end

# TODO: In the future this can be replaced by Msf::WebServices::HttpDBManagerService
def thin_cmd
  server_opts = "--rackup #{@ws_conf.shellescape} --address #{@options[:address].shellescape} --port #{@options[:port]}"
  ssl_opts = @options[:ssl] ? "--ssl --ssl-key-file #{@options[:ssl_key].shellescape} --ssl-cert-file #{@options[:ssl_cert].shellescape}" : ''
  ssl_opts << ' --ssl-disable-verify' if skip_ssl_verify?
  adapter_opts = "--environment #{@options[:ws_env]}"
  daemon_opts = "--daemonize --log #{@ws_log.shellescape} --pid #{@ws_pid.shellescape} --tag #{@ws_tag}" if @options[:daemon]
  all_opts = [server_opts, ssl_opts, adapter_opts, daemon_opts].reject(&:blank?).join(' ')

  "thin #{all_opts}"
end

def process_active?(pid)
  begin
    Process.kill(0, pid)
    true
  rescue Errno::ESRCH
    false
  end
end

def http_request(uri:, query: nil, data: nil, method: :get, headers: nil, skip_verify: false, cert: nil)
  all_headers = { 'User-Agent': @script_name }
  all_headers.merge!(headers) unless headers.nil?
  query_str = (!query.nil? && !query.empty?) ? URI.encode_www_form(query.compact) : nil
  uri.query = query_str

  http = Net::HTTP.new(uri.host, uri.port)
  if uri.is_a?(URI::HTTPS)
    http.use_ssl = true
    if skip_verify
      http.verify_mode = OpenSSL::SSL::VERIFY_NONE
    else
      # https://stackoverflow.com/questions/22093042/implementing-https-certificate-pubkey-pinning-with-ruby
      http.verify_mode = OpenSSL::SSL::VERIFY_PEER
      user_passed_cert = OpenSSL::X509::Certificate.new(File.read(cert))

      http.verify_callback = lambda do |preverify_ok, cert_store|
        server_cert = cert_store.chain[0]
        return true unless server_cert.to_der == cert_store.current_cert.to_der
        same_public_key?(server_cert, user_passed_cert)
      end
    end
  end

  begin
    response_data = { response: nil }
    case method
      when :get
        request = Net::HTTP::Get.new(uri.request_uri, initheader=all_headers)
      when :post
        request = Net::HTTP::Post.new(uri.request_uri, initheader=all_headers)
      else
        raise Exception, "Request method #{method} is not handled"
    end

    request.content_type = 'application/json'
    unless data.nil?
      json_body = data.to_json
      request.body = json_body
    end

    response = http.request(request)
    unless response.body.nil? || response.body.empty?
      response_data[:response] = JSON.parse(response.body, symbolize_names: true)
    end
  rescue => e
    response_data[:exception] = e
    puts "Problem with HTTP #{method} request #{uri.request_uri}, message: #{e.message}" if @options[:debug]
  end

  response_data
end

# Tells us whether the private keys on the passed certificates match
# and use the same algo
def same_public_key?(ref_cert, actual_cert)
  pkr, pka = ref_cert.public_key, actual_cert.public_key

  # First check if the public keys use the same crypto...
  return false unless pkr.class == pka.class
  # ...and then - that they have the same contents
  return false unless pkr.to_pem == pka.to_pem

  true
end

def parse_args(args)
  subtext = <<~USAGE
    Commands:
      init     initialize the component
      reinit   delete and reinitialize the component
      delete   delete and stop the component
      status   check component status
      start    start the component
      stop     stop the component
      restart  restart the component
  USAGE

  parser = OptionParser.new do |opts|
    opts.banner = "Usage: #{@script_name} [options] <command>"
    opts.separator('Manage a Metasploit Framework database and web service')
    opts.separator('')
    opts.separator('General Options:')
    opts.on('--component COMPONENT', @components + ['all'], 'Component used with provided command (default: database)',
            "  (#{@components.join(', ')})") { |component|
      @options[:component] = component.to_sym
    }

    opts.on('-d', '--debug', 'Enable debug output') { |d| @options[:debug] = d }
    opts.on('-h', '--help', 'Show this help message') {
      puts opts
      exit
    }
    opts.on('--use-defaults', 'Accept all defaults and do not prompt for options during an init') { |d|
      @options[:use_defaults] = d
    }

    opts.separator('')
    opts.separator('Database Options:')
    opts.on('--msf-db-name NAME', "Database name (default: #{@options[:msf_db_name]})") { |n|
      @options[:msf_db_name] = n
    }

    opts.on('--msf-db-user-name USER', "Database username (default: #{@options[:msf_db_user]})") { |u|
      @options[:msf_db_user] = u
    }

    opts.on('--msf-test-db-name NAME', "Test database name (default: #{@options[:msftest_db_name]})") { |n|
      @options[:msftest_db_name] = n
    }

    opts.on('--msf-test-db-user-name USER', "Test database username (default: #{@options[:msftest_db_user]})") { |u|
      @options[:msftest_db_user] = u
    }

    opts.on('--db-port PORT', Integer, "Database port (default: #{@options[:db_port]})") { |p|
      @options[:db_port] = p
    }

    opts.on('--db-pool MAX', Integer, "Database connection pool size (default: #{@options[:db_pool]})") { |m|
      @options[:db_pool] = m
    }

    opts.on('--connection-string URI', 'Use a pre-existing database cluster for initialization',
            'Example: --connection-string=postgresql://postgres:mysecretpassword@localhost:5432/postgres') { |c|
      @connection_string = c
    }

    opts.separator('')
    opts.separator('Web Service Options:')
    opts.on('-a', '--address ADDRESS',
            "Bind to host address (default: #{@options[:address]})") { |a|
      @options[:address] = a
    }

    opts.on('-p', '--port PORT', Integer,
            "Web service port (default: #{@options[:port]})") { |p|
      @options[:port] = p
    }

    opts.on('--[no-]daemon', 'Enable daemon') { |d|
      @options[:daemon] = d
    }

    opts.on('--[no-]ssl', "Enable SSL (default: #{@options[:ssl]})") { |s| @options[:ssl] = s }

    opts.on('--ssl-key-file PATH', "Path to private key (default: #{@options[:ssl_key]})") { |p|
      @options[:ssl_key] = p
    }

    opts.on('--ssl-cert-file PATH', "Path to certificate (default: #{@options[:ssl_cert]})") { |p|
      @options[:ssl_cert] = p
    }

    opts.on('--[no-]ssl-disable-verify',
            "Disables (optional) client cert requests (default: #{@options[:ssl_disable_verify]})") { |v|
      @options[:ssl_disable_verify] = v
    }

    opts.on('--environment ENV', @environments,
            "Web service framework environment (default: #{@options[:ws_env]})",
            "  (#{@environments.join(', ')})") { |e|
      @options[:ws_env] = e
    }

    opts.on('--retry-max MAX', Integer,
            "Maximum number of web service connect attempts (default: #{@options[:retry_max]})") { |m|
      @options[:retry_max] = m
    }

    opts.on('--retry-delay DELAY', Float,
            "Delay in seconds between web service connect attempts (default: #{@options[:retry_delay]})") { |d|
      @options[:retry_delay] = d
    }

    opts.on('--user USER', 'Initial web service admin username') { |u|
      @options[:ws_user] = u
    }

    opts.on('--pass PASS', 'Initial web service admin password') { |p|
      @options[:ws_pass] = p
    }

    opts.on('--[no-]msf-data-service NAME', 'Local msfconsole data service connection name') { |n|
      if !n
        @options[:add_data_service] = false
      else
        @options[:data_service_name] = n
      end
    }

    opts.separator('')
    opts.separator(subtext)
  end

  parser.parse!(args)

  if args.length != 1
    puts parser
    abort
  end

  @options
end

def invoke_command(commands, component, command)
  method = commands[component][command]
  if !method.nil?
    send(method)
  else
    print_error "Error: unrecognized command '#{command}' for #{component}"
  end
end

def installed?(cmd)
  !Msf::Util::Helper.which(cmd).nil?
end

def has_requirements(postgresql_cmds)
  ret_val = true
  other_cmds = %w(bundle thin)
  missing_msg = "Missing requirement: %<name>s does not appear to be installed or '%<prog>s' is not in the environment path"

  postgresql_cmds.each do |cmd|
    next unless Msf::Util::Helper.which(cmd).nil?
    puts missing_msg % { name: 'PostgreSQL', prog: cmd }
    ret_val = false
  end

  other_cmds.each do |cmd|
    if Msf::Util::Helper.which(cmd).nil?
      puts missing_msg % { name: "'#{cmd}'", prog: cmd }
      ret_val = false
    end
  end

  ret_val
end

def should_generate_web_service_ssl
  @options[:ssl] && ((!File.file?(@options[:ssl_key]) || !File.file?(@options[:ssl_cert])) ||
      (@options[:ssl_key] == @ws_ssl_key_default && @options[:ssl_cert] == @ws_ssl_cert_default))
end

def prompt_for_component(command)
  if command == :status || command == :delete
    return :all
  end

  if command == :stop && web_service_pid_status != WebServicePIDStatus::RUNNING
    return :database
  end

  enable_webservice = ask_yn("Would you like to #{command} the webservice? (Not Required)", default: 'no')
  if enable_webservice
    :all
  else
    :database
  end
end

def prompt_for_deletion(command)
  destructive_operations = [:init, :reinit, :delete]

  if destructive_operations.include? command
    if command == :init
      return if web_service_pid_status != WebServicePIDStatus::NO_PID_FILE
      if (@options[:component] == :all || @options[:component] == :webservice) && should_generate_web_service_ssl &&
          (File.file?(@options[:ssl_key]) || File.file?(@options[:ssl_cert]))
        @options[:delete_existing_data] = should_delete
        return
      end
      if (@options[:component] == :all || @options[:component] == :database) && File.exist?(@db_conf)
        @options[:delete_existing_data] = should_delete
        return
      end
    else
      @options[:delete_existing_data] = should_delete
    end
  end
end

def should_delete
  return true if @options[:use_defaults]
  ask_yn("Would you like to delete your existing data and configurations?")
end

if $PROGRAM_NAME == __FILE__
  # Bomb out if we're root
  if !Gem.win_platform? && Process.uid.zero?
    puts "Please run #{@script_name} as a non-root user"
    abort
  end

  # map component commands to methods
  commands = {
      database: {
          init: :init_db,
          reinit: :reinit_db,
          delete: :delete_db,
          status: :status_db,
          start: :start_db,
          stop: :stop_db,
          restart: :restart_db
      },
      webservice: {
          init: :init_web_service,
          reinit: :reinit_web_service,
          delete: :delete_web_service,
          status: :status_web_service,
          start: :start_web_service,
          stop: :stop_web_service,
          restart: :restart_web_service
      }
  }

  parse_args(ARGV)
  update_db_port

  if @connection_string
    @db_driver = MsfdbHelpers::Standalone.new(options: @options, db_conf: @db_conf, connection_string: @connection_string)
  elsif installed?('pg_ctl') && has_requirements(MsfdbHelpers::PgCtl.requirements)
    @db_driver = MsfdbHelpers::PgCtl.new(db_path: @db, options: @options, localconf: @localconf, db_conf: @db_conf)
  elsif installed?('pg_ctlcluster') && has_requirements(MsfdbHelpers::PgCtlcluster.requirements)
    @db_driver = MsfdbHelpers::PgCtlcluster.new(db_path: @db, options: @options, localconf: @localconf, db_conf: @db_conf)
  else
    print_error('You need to have postgres installed or specify a database with --connection-string')
    abort
  end

  command = ARGV[0].to_sym
  if @options[:component].nil?
    @options[:component] = prompt_for_component(command)
  end
  prompt_for_deletion(command)
  # Disable the locally configured webservice component
  if command != :status && @options[:component] == :database
    clear_default_data_service
  end
  if @options[:component] == :all
    @components.each { |component|
      puts '===================================================================='
      puts "Running the '#{command}' command for the #{component}:"
      invoke_command(commands, component.to_sym, command)
      puts '===================================================================='
      puts
    }
  else
    puts "Running the '#{command}' command for the #{@options[:component]}:"
    invoke_command(commands, @options[:component], command)
  end
end