#! /usr/bin/env ruby

def mdm_banner
  banner = {}
  banner[:color] = "\e[34m"
  banner[:text] = <<-BANNER
 _______ _______________________ _______ _______ _       _______ __________________
(       )  ____ \\__   __/  ___  )  ____ \\  ____ ) \\     (  ___  )\\__   __/\\__   __/
| () () | (    \\/  ) (  | (   ) | (    \\/ (    )| (     | (   ) |   ) (      ) (   
| || || | (__      | |  | (___) | (_____| (____)| |     | |   | |   | |      | |   
| |(_)| |  __)     | |  |  ___  |_____  )  _____) |     | |   | |   | |      | |   
| |   | | (        | |  | (   ) |     ) | (     | |     | |   | |   | |      | |   
| )   ( | (____/\\  | |  | )   ( |\\____) | )     | (____/\\ (___) |___) (___   | |   
|/     \\|_______/  )_(  |/     \\|_______)/      (_______/_______)\\_______/   )_(   
                                                                                   

 ______  _______________________ _______ _______ ______  _______ _       _______ 
(  __  \\(  ___  )__   __/  ___  )       )  ___  )  __  \\(  ____ \\ \\     (  ____ \\
| (  \\  ) (   ) |  ) (  | (   ) | () () | (   ) | (  \\  ) (    \\/ (     | (    \\/
| |   ) | (___) |  | |  | (___) | || || | |   | | |   ) | (__   | |     | (_____ 
| |   | |  ___  |  | |  |  ___  | |(_)| | |   | | |   | |  __)  | |     (_____  )
| |   ) | (   ) |  | |  | (   ) | |   | | |   | | |   ) | (     | |           ) |
| (__/  ) )   ( |  | |  | )   ( | )   ( | (___) | (__/  ) (____/\\ (____/Y\\____) |
(______/|/     \\|  )_(  |/     \\|/     \\|_______)______/(_______/_______|_______)
BANNER
  banner
end

def db_info_file
  hidden_file = "#{Dir.home}/.mdm.yml"
  if File.readable?(hidden_file)
    pro_path = YAML.load_file(hidden_file)['pro_path']
    return "#{pro_path}/ui/config/database.yml"
  elsif !ARGV[0].blank?
    return ARGV[0]
  else
    warn "No YAML file of DB info available"
    exit
  end
end

begin
  require 'pry'
  require "metasploit_data_models"

  # Set up a DB connection, preferring one from Pro if it's in the normal place
  # Otherwise get from ARGV[0]

  if File.readable?(db_info_file)
    connection_info = YAML.load_file(db_info_file)
    ActiveRecord::Base.establish_connection(connection_info['development'])
  else
    warn "Can't access DB -- check file path."
    exit
  end

  MetasploitDataModels.require_models

  puts "\n\n\n#{mdm_banner[:color]}#{mdm_banner[:text]}\e[0m\n\n\n"

  Pry.config.prompt = proc { |obj, nest_level, _| "mdm:#{nest_level}> " }

  Pry.start
  exit
rescue LoadError
  warn "Unable to load Pry"
end