#!/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.config_directory
@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: false,
    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
    raise MsfdbHelpers::SecurityError, 'Database startup failed. Inspect the database log; do not reinitialize or delete data to repair authentication.'
  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'
    raise MsfdbHelpers::SecurityError, 'Could not start the existing database' unless @db_driver.start

    return
  end

  if @db_driver.exists? && !@options[:delete_existing_data] && !load_db_config
    raise MsfdbHelpers::SecurityError, 'Could not load existing database credentials; repair the configuration without reinitializing'
  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
    unless @db_driver.run_cmd('bundle exec rake db:migrate').zero?
      raise MsfdbHelpers::SecurityError, 'Database initialization did not complete: schema migration failed. Inspect the migration error without reinitializing the database.'
    end
  end
  puts 'Database initialization successful'.green.bold.to_s
end

def load_db_config
  if File.file?(@db_conf)
    config = load_db_config_yaml

    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', 0o600) do |f|
    f.chmod(0o600) unless Gem.win_platform?
    f.puts <<~EOF
      development: &pgsql
        adapter: postgresql
        database: #{@options[:msf_db_name].to_json}
        username: #{@options[:msf_db_user].to_json}
        password: #{@msf_pass.to_json}
        host: #{@options[:db_host].to_json}
        port: #{@options[:db_port]}
        pool: #{@options[:db_pool]}

      production: &production
        <<: *pgsql

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

# Loads the database config, tolerating Psych versions that do not accept the aliases keyword.
#
# @return [Hash] the parsed config, or an empty Hash if the file is empty.
def load_db_config_yaml
  YAML.load_file(@db_conf, aliases: true) || {}
rescue ArgumentError
  YAML.load_file(@db_conf) || {}
end

def update_db_port
  if File.file?(@db_conf)
    config = load_db_config_yaml
    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

def print_webservice_removal_prompt
  $stderr.puts "#{'[WARNING]'.red} The remote web service is being removed. Does this impact you? React here: https://github.com/rapid7/metasploit-framework/issues/18439"
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

  # The first account is created directly in the database before the service starts, so the web
  # service never has to accept an unauthenticated request to create it.
  if add_web_service_user && start_web_service && add_web_service_workspace
    output_web_service_information
  else
    puts 'Failed to complete MSF web service configuration, please reinitialize.'
    stop_web_service
  end
end

def start_web_service_daemon
  if @db_driver.run_cmd("#{thin_cmd} start") == 0
    # wait until web service is online
    retry_count = 0
    response_data = web_service_online_check
    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
      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
  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
  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
  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
    # The service always enforces authentication, so an unauthenticated request is expected to be rejected
    response_data[:state] = :online
  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: 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,
                               headers: { 'Authorization' => "Bearer #{@ws_api_token}" },
                               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}"

  # Create the admin user and its API token directly in the database, rather than over HTTP, so that
  # no unauthenticated request to the web service is ever able to create the first account.
  result = create_web_service_user_token
  if result[:token].nil?
    print_error "Error creating MSF web service user #{@msf_ws_user}: #{result[:error]}"
    return false
  end
  @ws_api_token = result[:token]

  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}"
  puts "#{'MSF web service user API token'.cyan.bold}: #{@ws_api_token}"
  return true
end

# Creates the initial admin user and issues its API token using the database that the web service
# is configured to use.
#
# @return [Hash] { token: String } on success; { token: nil, error: String } describing why the user
#   could not be created otherwise.
def create_web_service_user_token
  config = load_db_config_yaml['production']
  return { token: nil, error: "No production section found in database config #{@db_conf}" } if config.nil?

  # Only the database is needed, so skip loading modules and file logging.
  framework = Msf::Simple::Framework.create('SkipDatabaseInit' => true, 'DeferModuleLoads' => true, 'DisableLogging' => true)
  unless framework.db.connect(config) && framework.db.active
    return { token: nil, error: "Could not connect to the database configured in #{@db_conf}: #{framework.db.error}" }
  end

  user = framework.db.report_user(username: @msf_ws_user, password: @msf_ws_pass, admin: true)
  { token: framework.db.create_new_user_token(id: user.id, token_length: 40) }
rescue StandardError => e
  { token: nil, error: "#{e.class}: #{e.message}" }
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
  # @msf_command holds the initial common part of commands (msfconsole -qx) and takes the optional specific commands as arguments (#{cmd})
  msf_command = "msfconsole -qx '#{cmd}'"
  if @db_driver.run_cmd(msf_command) != 0
    # attempt to execute msfconsole in the current working directory
    if @db_driver.run_cmd(msf_command, 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 = "#{get_db_connect_command}; db_save; 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[:add_data_service] = true
        @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
rescue MsfdbHelpers::SecurityError, SystemCallError => e
  raise unless component == :database

  warn "Database operation failed: #{e.message}"
  exit 1
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

  if @options[:add_data_service] == true
    :all
  else
    :database
  end
end

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

  if destructive_operations.include? command
    @options[:delete_existing_data] = should_delete
  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 File.expand_path($PROGRAM_NAME) == File.expand_path(__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)
  if @options[:component] == :all
    @components.each { |component|
      if component == :webservice
        3.times { print_webservice_removal_prompt }
      end
      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]}:"
    if @options[:component] == :webservice
      3.times { print_webservice_removal_prompt }
    end
    invoke_command(commands, @options[:component], command)
  end
end