#!/usr/bin/env ruby
# -*- coding: binary -*-
#
# $Id$
#
# This keeps the framework up-to-date
#
# $Revision$
#

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

@orig_dir = Dir.pwd
@msfbase_dir = File.dirname(msfbase)

Dir.chdir(@msfbase_dir)

def print_usage
  puts <<-EOF
usage: msfupdate [--git-remote=<remote>] [--git-branch=<branch>] [--offline-file=<file>]
Options:
-h, --help               show help
    --git-remote REMOTE  git remote to use (default upstream)
    --git-branch BRANCH  git branch to use (default master)
    --offline-file FILE  offline update file to use
    EOF
end

require 'getoptlong'
opts = GetoptLong.new(
  ['--help', '-h', GetoptLong::NO_ARGUMENT],
  ['--git-remote', GetoptLong::REQUIRED_ARGUMENT],
  ['--git-branch', GetoptLong::REQUIRED_ARGUMENT],
  ['--offline-file', GetoptLong::REQUIRED_ARGUMENT],
)

begin
  opts.each do |opt, arg|
    case opt
    when '--help'
      print_usage
      exit
    when '--git-remote'
      @git_remote = arg
    when '--git-branch'
      @git_branch = arg
    when '--offline-file'
      @offline_file = File.expand_file(arg, @orig_dir)
    end
  end
rescue GetoptLong::Error
  $stderr.puts "#{$0}: try 'msfupdate --help' for more information"
  exit 0x20
end

# Handle the old wait/nowait argument behavior
if ARGV[0] == 'wait' || ARGV[0] == 'nowait'
  @actually_wait = (ARGV.shift == 'wait')
end

$stderr.puts "[*]"
$stderr.puts "[*] Attempting to update the Metasploit Framework..."
$stderr.puts "[*]"
$stderr.puts ""

# Bail right away, no waiting around for consoles.
if not (Process.uid == 0 or File.stat(msfbase).owned?)
  $stderr.puts "[-] ERROR: User running msfupdate does not own the Metasploit installation"
  $stderr.puts "[-] Please run msfupdate as the same user who installed Metasploit."
  exit 0x10
end

def is_apt
  File.exists?(File.expand_path(File.join(@msfbase_dir, '.apt')))
end

# Are you an installer, or did you get here via a source checkout?
def is_installed
  File.exists?(File.expand_path(File.join(@msfbase_dir, "..", "engine", "update.rb"))) && !is_apt
end

def is_git
  File.directory?(File.join(@msfbase_dir, ".git"))
end

# Adding an upstream enables msfupdate to pull updates from
# Rapid7's metasploit-framework repo instead of the repo
# the user originally cloned or forked.
def add_git_upstream
  $stdout.puts "[*] Attempting to add remote 'upstream' to your local git repository."
  system("git", "remote", "add", "upstream", "git://github.com/rapid7/metasploit-framework.git")
  $stdout.puts "[*] Added remote 'upstream' to your local git repository."
end

# This only exits if you actually pass a wait option, otherwise
# just returns nil. This is likely unexpected, revisit this.
def maybe_wait_and_exit(exit_code=0)
  if @actually_wait
    $stdout.puts ""
    $stdout.puts "[*] Please hit enter to exit"
    $stdout.puts ""
    $stdin.readline
    exit exit_code
  end
end

def apt_upgrade_available(package)
  require 'open3'
  installed = nil
  upgrade = nil
  ::Open3.popen3({'LANG'=>'en_US.UTF-8'}, "apt-cache", "policy", package) do |stdin, stdout, stderr|
    stdout.each do |line|
      installed = $1 if line =~ /Installed: ([\w\-+.:~]+)$/
      upgrade = $1 if line =~ /Candidate: ([\w\-+.:~]+)$/
      break if installed && upgrade
    end
  end
  if installed && installed != upgrade
    upgrade
  else
    nil
  end
end

####### Since we're Git, do it all that way #######
if is_git
  $stdout.puts "[*] Checking for updates via git"
  $stdout.puts "[*] Note: Updating from bleeding edge"
  out = `git remote show upstream` # Actually need the output for this one.
  add_git_upstream unless $?.success? and out =~ %r{(https|git|git@github\.com):(//github\.com/)?(rapid7/metasploit-framework\.git)}

  remote = @git_remote || "upstream"
  branch = @git_branch || "master"

  # This will save local changes in a stash, but won't
  # attempt to reapply them. If the user wants them back
  # they can always git stash pop them, and that presumes
  # they know what they're doing when they're editing local
  # checkout, which presumes they're not using msfupdate
  # to begin with.
  #
  # Note, this requires at least user.name and user.email
  # to be configured in the global git config. Installers should
  # take care that this is done. TODO: Enforce this in msfupdate
  committed = system("git", "diff", "--quiet", "HEAD")
  if committed.nil?
    $stderr.puts "[-] ERROR: Failed to run git"
    $stderr.puts ""
    $stderr.puts "[-] If you used a binary installer, make sure you run the symlink in"
    $stderr.puts "[-] /usr/local/bin instead of running this file directly (e.g.: ./msfupdate)"
    $stderr.puts "[-] to ensure a proper environment."
    maybe_wait_and_exit 1
  elsif not committed
    system("git", "stash")
    $stdout.puts "[*] Stashed local changes to avoid merge conflicts."
    $stdout.puts "[*] Run `git stash pop` to reapply local changes."
  end

  system("git", "reset", "HEAD", "--hard")
  system("git", "checkout", branch)
  system("git", "fetch", remote)
  system("git", "merge", "#{remote}/#{branch}")

  $stdout.puts "[*] Updating gems..."
  require 'bundler'
  Bundler.with_clean_env do
    system("bundle", "install")
  end
end

if is_installed
  update_script = File.expand_path(File.join(@msfbase_dir, "..", "engine", "update.rb"))
  product_key =   File.expand_path(File.join(@msfbase_dir, "..", "engine", "license", "product.key"))
  if File.exists? product_key
    if File.readable? product_key
      if (@offline_file)
        system("ruby", update_script, @offline_file)
      else
        system("ruby", update_script)
      end
    else
      $stdout.puts "[-] ERROR: Failed to update Metasploit installation"
      $stdout.puts ""
      $stdout.puts "[-] You must be able to read the product key for the"
      $stdout.puts "[-]	Metasploit installation in order to run msfupdate."
      $stdout.puts "[-] Usually, this means you must be root (EUID 0)."
      maybe_wait_and_exit 10
    end
  else
    $stdout.puts "[-] ERROR: Failed to update Metasploit installation"
    $stdout.puts ""
    $stdout.puts "[-] In order to update your Metasploit installation,"
    $stdout.puts "[-] you must first register it through the UI, here:"
    $stderr.puts "[-] https://localhost:3790"
	$stderr.puts "[-] (Note: Metasploit Community Edition is totally"
	$stderr.puts "[-] free and takes just a few seconds to register!)"
    maybe_wait_and_exit 11
  end
end

if is_apt
  # For more information, see here:
  #   	https://community.rapid7.com/community/metasploit/blog/2013/01/17/metasploit-updates-and-msfupdate
  $stdout.puts "[*] Checking for updates via the APT repository"
  $stdout.puts "[*] Note: expect weekly(ish) updates using this method"
  system("apt-get", "-qq", "update")

  packages = []
  packages << 'metasploit-framework' if framework_version = apt_upgrade_available('metasploit-framework')
  packages << 'metasploit' if pro_version = apt_upgrade_available('metasploit')

  if packages.empty?
    $stdout.puts "[*] No updates available"
  else
    $stdout.puts "[*] Updating to version #{pro_version || framework_version}"
    system("apt-get", "install", "--assume-yes", *packages)
    if packages.include?('metasploit')
      start_cmd = File.expand_path(File.join(@msfbase_dir, '..', '..', '..', 'scripts', 'start.sh'))
      system(start_cmd) if ::File.executable_real? start_cmd
    end
  end
end

unless is_git || is_installed || is_apt
  raise RuntimeError, "Cannot determine checkout type: `#{@msfbase_dir}'"
end

maybe_wait_and_exit(0)