Update msfupdate to accommodate for forks and clones.

Author: erran

msfupdate

@@ -40,6 +40,15 @@ 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/rapid7metasploit-framework.git")
+	$stdout.puts "[*] Added remote 'upstream' to your local git repository."
+end
+
 def print_deprecation_warning
 	$stdout.puts "[*] Deprecation Note: The next version of Metasploit will"
 	$stdout.puts "[*] update over the git protocol, which requires outbound"
@@ -102,8 +111,12 @@ end
 
 ####### Since we're Git, do it all that way #######
 if is_git
-	remote = @git_remote || "origin"
-	branch = @git_branch || "master"
+	out = `git remote show upstream`
+	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_branchch || "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
@@ -114,22 +127,23 @@ if is_git
 	# 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
-	res = system("git", "stash")
-	if res.nil?
+	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."
 		exit 1
-	else
-		$stdout.puts "[*] Stashed local changes (if any) to avoid merge conflicts."
-		$stdout.puts "[*] Run 'git stash pop' to reapply local changes."
+	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")
+	system("git", "fetch", remote)
 	system("git", "merge", "#{remote}/#{branch}")
 end