#!/usr/bin/env ruby

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

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

@localconf = "#{ENV['HOME']}/.msf4"
@db = "#{@localconf}/db"
@db_conf = "#{@localconf}/database.yml"

@ws_tag = 'msf-ws'
@ws_conf = "#{@localconf}/#{@ws_tag}-config.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]

@options = {
    component: :all,
    debug: false,
    msf_db_name: 'msf',
    msf_db_user: 'msf',
    msftest_db_name: 'msftest',
    msftest_db_user: 'msftest',
    db_port: 5433,
    db_pool: 200,
    address: 'localhost',
    port: 8080,
    ssl: true,
    ssl_cert: @ws_ssl_cert_default,
    ssl_key: @ws_ssl_key_default,
    ssl_disable_verify: false,
    ws_env: ENV['RACK_ENV'] || 'production',
    retry_max: 10,
    retry_delay: 5.0,
    ws_user: nil
}


def run_cmd(cmd, input = nil)
  exitstatus = 0
  err = out = ""

  puts "run_cmd: cmd=#{cmd}, input=#{input}" if @options[:debug]

  Open3.popen3(cmd) do |stdin, stdout, stderr, wait_thr|
    stdin.puts(input) if input
    if @options[:debug]
      err = stderr.read
      out = stdout.read
    end
    exitstatus = wait_thr.value.exitstatus
  end

  if exitstatus != 0
    if @options[:debug]
      puts "'#{cmd}' returned #{exitstatus}"
      puts out
      puts err
    end
  end

  exitstatus
end

def run_psql(cmd, db_name: 'postgres')
  if @options[:debug]
    puts "psql -p #{@options[:db_port]} -c \"#{cmd};\" #{db_name}"
  end

  run_cmd("psql -p #{@options[:db_port]} -c \"#{cmd};\" #{db_name}")
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

  if Dir.exist?(@db)
    if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db} status") == 0
      puts "Database started at #{@db}"
    else
      puts "Database is not running at #{@db}"
    end
  else
    puts "No database found at #{@db}"
  end
end

def started_db
  if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db} status") == 0
    puts "Database already started at #{@db}"
    return true
  end

  print "Starting database at #{@db}..."
  run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db} -l #{@db}/log start")
  sleep(2)
  if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db} status") != 0
    puts 'failed'
    false
  else
    puts 'success'
    true
  end
end

def start_db
  if !Dir.exist?(@db)
    puts "No database found at #{@db}, not starting"
    return
  end

  update_db_port

  while !started_db
    last_log = tail("#{@db}/log")
    puts last_log
    fixed = false
    if last_log =~ /not compatible/
      puts "Please attempt to upgrade the database manually using pg_upgrade."
    end
    if !fixed
      if ask_yn('If your database is corrupt, would you to reinitialize it?')
        fixed = reinit_db
      end
    end
    if !fixed
      if !ask_yn('Database not started, try again?')
        return
      end
    end
  end
end

def stop_db
  update_db_port

  if run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db} status") == 0
    puts "Stopping database at #{@db}"
    run_cmd("pg_ctl -o \"-p #{@options[:db_port]}\" -D #{@db} stop")
  else
    puts "Database is no longer running at #{@db}"
  end
end

def restart_db
  stop_db
  start_db
end

def create_db
  puts "Creating database at #{@db}"
  Dir.mkdir(@db)
  run_cmd("initdb --auth-host=trust --auth-local=trust -E UTF8 #{@db}")

  File.open("#{@db}/pg_hba.conf", 'w') do |f|
    f.puts "host    \"msf\"      \"#{@options[:msf_db_user]}\"      127.0.0.1/32           md5"
    f.puts "host    \"msftest\"  \"#{@options[:msftest_db_user]}\"  127.0.0.1/32           md5"
    f.puts "host    \"postgres\"  \"#{@options[:msftest_db_user]}\"  127.0.0.1/32           trust"
    f.puts "host    \"template1\"   all                127.0.0.1/32           trust"
    if Gem.win_platform?
      f.puts "host    all             all                127.0.0.1/32           trust"
      f.puts "host    all             all                ::1/128                trust"
    else
      f.puts "local   all             all                                       trust"
    end
  end

  File.open("#{@db}/postgresql.conf", 'a') do |f|
    f.puts "port = #{@options[:db_port]}"
  end
end

def init_db
  if Dir.exist?(@db)
    puts "Found a database at #{@db}, checking to see if it is started"
    start_db
    return
  end

  if File.exist?(@db_conf)
    if !ask_yn("Found database config at #{@db_conf}, do you want to overwrite it?")
      return
    end
  end

  # Generate new database passwords
  msf_pass = pw_gen
  msftest_pass = pw_gen

  # 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: 127.0.0.1
        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)

  create_db

  start_db

  puts 'Creating database users'
  run_psql("create user #{@options[:msf_db_user]} with password '#{msf_pass}'")
  run_psql("create user #{@options[:msftest_db_user]} with password '#{msftest_pass}'")
  run_psql("alter role #{@options[:msf_db_user]} createdb")
  run_psql("alter role #{@options[:msftest_db_user]} createdb")
  run_psql("alter role #{@options[:msf_db_user]} with password '#{msf_pass}'")
  run_psql("alter role #{@options[:msftest_db_user]} with password '#{msftest_pass}'")
  run_cmd("createdb -p #{@options[:db_port]} -O #{@options[:msf_db_user]} -h 127.0.0.1 -U #{@options[:msf_db_user]} -E UTF-8 -T template0 #{@options[:msf_db_name]}",
    "#{msf_pass}\n#{msf_pass}\n")
  run_cmd("createdb -p #{@options[:db_port]} -O #{@options[:msftest_db_user]} -h 127.0.0.1 -U #{@options[:msftest_db_user]} -E UTF-8 -T template0 #{@options[:msftest_db_name]}",
    "#{msftest_pass}\n#{msftest_pass}\n")

  puts 'Creating initial database schema'
  Dir.chdir(@framework) do
    run_cmd("bundle exec rake db:migrate")
  end
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)
  loop do
    print "#{question}: "
    yn = STDIN.gets
    case yn
    when /^[Yy]/
      return true
    when /^[Nn]/
      return false
    else
      puts 'Please answer yes or no.'
    end
  end
end

def ask_value(question, default_value)
  print "#{question}[#{default_value}]: "
  input = STDIN.gets.strip
  if input.nil? || input.empty?
    return default_value
  else
    return input
  end
end


def delete_db
  if Dir.exist?(@db)
    puts "Deleting all data at #{@db}"
    stop_db
    FileUtils.rm_rf(@db)
    if File.exist?(@db_conf) && ask_yn("Delete database configuration at #{@db_conf}?")
      File.delete(@db_conf)
    end
  else
    puts "No data at #{@db}, doing nothing"
  end
end

def reinit_db
  delete_db
  init_db
end

def status_web_service
  if File.file?(@ws_pid)
    ws_pid = tail(@ws_pid)
    if ws_pid.nil? || !process_active?(ws_pid.to_i)
      puts "MSF web service PID file found, but no active process running as PID #{ws_pid}"
    else
      puts "MSF web service is running as PID #{ws_pid}"
    end
  else
    puts "No MSF web service PID file found at #{@ws_pid}"
  end
end

def init_web_service
  if File.file?(@ws_conf)
    if !ask_yn("Found web service config at #{@ws_conf}, do you want to overwrite it?")
      return
    end
  end

  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

  # Write a default Rack config file for the web service
  Dir.mkdir(@localconf) unless File.directory?(@localconf)
  File.open(@ws_conf, 'w') do |f|
    f.puts <<~EOF
      # #{File.basename(@ws_conf)}
      # created on: #{Time.now.utc}
      lib_path = File.expand_path('./lib/', '#{@framework}')
      $LOAD_PATH << lib_path unless $LOAD_PATH.include?(lib_path)
      require 'msf/core/db_manager/http/metasploit_api_app'
      run MetasploitApiApp
    EOF
  end
  File.chmod(0640, @ws_conf)

  if @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))
    generate_web_service_ssl(key: @options[:ssl_key], cert: @options[:ssl_cert])
  end

  if !start_web_service
    return
  end

  # wait until web service is online
  retry_count = 0
  is_online = web_service_online
  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])
    is_online = web_service_online
  end

  if is_online
    add_web_service_user
  else
    puts "MSF web service does not appear to be online; aborting initialize."
  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

  # daemonize MSF web service
  puts "Starting MSF web service"
  if run_cmd("#{thin_cmd} start") == 0
    puts "MSF web service started"
    return true
  else
    puts "MSF web service not started"
    return false
  end
end

def stop_web_service
  ws_pid = tail(@ws_pid)
  if ws_pid.nil? || !process_active?(ws_pid.to_i)
    puts "MSF web service is no longer running"
    if File.file?(@ws_pid)
      puts "Deleting MSF web service PID file #{@ws_pid}"
      File.delete(@ws_pid)
    end
  else
    puts "Stopping MSF web service PID #{ws_pid}"
    run_cmd("#{thin_cmd} stop")
  end
end

def restart_web_service
  stop_web_service
  start_web_service
end

def delete_web_service
  stop_web_service
  if File.file?(@ws_conf) && ask_yn("Delete MSF web service configuration at #{@ws_conf}?")
    File.delete(@ws_conf)
  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)) &&
      !ask_yn("Either MSF web service SSL key #{key} or certificate #{cert} already exist, overwrite both?")
    return
  end

  puts "Generating SSL key and certificate for MSF web service"
  # @ssl_cert = Rex::Socket::SslTcpServer.ssl_generate_certificate
  @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
  msf_version_uri = get_web_service_uri(path: '/api/v1/msf/version')
  response = http_request(uri: msf_version_uri, method: :get,
                          skip_verify: skip_ssl_verify, cert: get_ssl_cert)
  puts "web_service_online: response=#{response}" if @options[:debug]
  !response.nil? && !response.dig(:data, :metasploit_version).nil?
end

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

  # Generate new web service user password
  msf_ws_pass = pw_gen
  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 = http_request(uri: user_uri, data: user_data, method: :post,
                          skip_verify: skip_ssl_verify, cert: get_ssl_cert)
  puts "add_web_service_user: create user response=#{response}" if @options[:debug]
  if response.nil? || response.dig(:data, :username) != @msf_ws_user
    puts "Error creating MSF web service user #{@msf_ws_user}"
    return false
  end
  puts "\nMSF web service user: #{@msf_ws_user}"
  puts "MSF web service password: #{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 = http_request(uri: generate_token_uri, query: cred_data, method: :get,
                          skip_verify: skip_ssl_verify, cert: get_ssl_cert)
  puts "add_web_service_user: generate token response=#{response}" if @options[:debug]
  if response.nil? || (@ws_api_token = response.dig(:data, :token)).nil?
    puts "Error creating MSF web service user API token"
    return false
  end
  puts "MSF web service user API token: #{@ws_api_token}"
  puts "Please store these credentials securely."
  return true
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

def thin_cmd
  server_opts = "--rackup #{@ws_conf} --address #{@options[:address]} --port #{@options[:port]}"
  ssl_opts = @options[:ssl] ? "--ssl --ssl-key-file #{@options[:ssl_key]} --ssl-cert-file #{@options[:ssl_cert]}" : ''
  ssl_opts << ' --ssl-disable-verify' if skip_ssl_verify
  adapter_opts = "--environment #{@options[:ws_env]}"
  daemon_opts = "--daemonize --log #{@ws_log} --pid #{@ws_pid} --tag #{@ws_tag}"
  all_opts = [server_opts, ssl_opts, adapter_opts, daemon_opts].reject(&:empty?).join(' ')

  "thin #{all_opts}"
end

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

def http_request(uri:, query: nil, data: nil, method: :get, skip_verify: false, cert: nil)
  headers = { 'User-Agent': @script_name }
  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
    case method
      when :get
        request = Net::HTTP::Get.new(uri.request_uri, initheader=headers)
      when :post
        request = Net::HTTP::Post.new(uri.request_uri, initheader=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?
      return JSON.parse(response.body, symbolize_names: true)
    end
  rescue EOFError => e
    puts "No data was returned for HTTP #{method} request #{uri.request_uri}, message: #{e.message}" if @options[:debug]
  rescue => e
    puts "Problem with HTTP #{method} request #{uri.request_uri}, message: #{e.message}" if @options[:debug]
  end
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: all)",
            "  (#{@components.join(', ')})") { |component|
      @options[:component] = component.to_sym
    }

    opts.on("-d", "--debug", "Enable debug output") { |d| @options[:debug] = 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.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-]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", ['production', 'development'],
            "Web service framework environment (default: #{@options[:ws_env]})",
            "  (production, development)") { |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 (default: #{@options[:ws_user]})") { |u|
      @options[:ws_user] = u
    }

    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
    puts "Error: unrecognized command '#{command}' for #{component}"
  end
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)

  command = ARGV[0].to_sym
  if @options[:component] == :all
    @components.each { |component|
      invoke_command(commands, component.to_sym, command)
    }
  else
    invoke_command(commands, @options[:component], command)
  end
end