#!/usr/bin/env ruby
# -*- coding: binary -*-

# $Id$
# $Revision$

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

@msfbase_dir = File.dirname(msfbase)

@args = ARGV.dup

# May be changed
@configdir = File.expand_path(File.join(File.dirname(msfbase), "data", "svn"))

Dir.chdir(@msfbase_dir)

$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

def is_svn
	File.directory?(File.join(@msfbase_dir, ".svn"))
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

def print_deprecation_warning
	$stdout.puts ""
	$stdout.puts "[-] Deprecation Note: Metasploit source checkouts NO LONGER update"
	$stdout.puts "[-] over SVN. You will need to reinstall Metasploit using"
	$stdout.puts "[-] binary installers (from http://www.metasploit.com/download ),"
	$stdout.puts "[-] Debian packages (currently only supported on Kali Linux), or"
	$stdout.puts "[-] a development source checkout from GitHub (see http://r-7.co/ZLhA8P )"
	$stdout.puts "[-] "
	$stdout.puts "[-] For more on msfupdate and migrating off of SVN, see http://r-7.co/MSF-UP"
	$stdout.puts ""
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

# Some of these args are meaningful for SVN, some for Git,
# some for both. Fun times.
@args.each_with_index do |arg,i|
	case arg
		# Handle the old wait/nowait argument behavior
	when "wait", "nowait"
		@wait_index = i
		@actually_wait = (arg == "wait")
		# An empty or absent config-dir means a default config-dir
	when "--config-dir"
		@configdir_index = i
		# A defined config dir means a defined config-dir
	when /--config-dir=(.*)?/
		# Spaces in the directory should be fine since this whole thing is passed
		# as a single argument via the multi-arg syntax for system() below.
		@configdir = $1
		@configdir_index = i
	when /--git-remote=([^\s]*)?/
		@git_remote = $1
		@git_remote_index = i
	when /--git-branch=([^\s]*)?/
		@git_branch = $1
		@git_branch_index = i
	end
end

@args[@wait_index] = nil      if @wait_index
@args[@configdir_index] = nil if @configdir_index

@args[@git_remote_index] = nil if @git_remote_index
@args[@git_branch_index] = nil if @git_branch_index
@args = @args.compact

####### Since we're SVN, do it all this way #######
if is_svn
	# We're fully deprecated now, so just exit.
	# Leaving in the commented code in case someone wants to
	# get a last-chance at msfupdate before the SVN server goes
	# off line, which will be ANY DAY NOW. Seriously.
	print_deprecation_warning
	$stdin.readline if @actually_wait
	exit(0x11) # Comment this to get old functionality back.
	@args.push("--config-dir=#{@configdir}")
	@args.push("--non-interactive")

	res = system("svn", "cleanup")
	if res.nil?
		$stderr.puts "[-] ERROR: Failed to run svn"
		$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
	else
		# Cleanup worked, go ahead and update
		system("svn", "update", *@args)
	end
end

####### Since we're Git, do it all that way #######
if is_git
	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
			system("ruby", update_script)
		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 (note, Metasploit Community"
		$stderr.puts "[-] Edition is totally free and takes just a few seconds"
		$stderr.puts "[-] to register!)"
		maybe_wait_and_exit 11
	end
end

if is_apt
	$stdout.puts "[*] Checking for updates"
	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_svn || is_git || is_installed || is_apt
	raise RuntimeError, "Cannot determine checkout type: `#{@msfbase_dir}'"
end

maybe_wait_and_exit(0)