diff --git a/lib/updater/git_updater.rb b/lib/updater/git_updater.rb index 0cfdecac..0afb4336 100644 --- a/lib/updater/git_updater.rb +++ b/lib/updater/git_updater.rb @@ -34,6 +34,14 @@ class GitUpdater < Updater %x[git #{repo_directory_arguments()} pull] end + def has_local_changes? + %x[git #{repo_directory_arguments()} diff --exit-code 2>&1] =~ /diff/ ? true : false + end + + def reset_head + %x[git #{repo_directory_arguments()} reset --hard HEAD] + end + protected def repo_directory_arguments if @repo_directory diff --git a/spec/lib/updater/git_updater_spec.rb b/spec/lib/updater/git_updater_spec.rb index 91293642..153cd0cc 100644 --- a/spec/lib/updater/git_updater_spec.rb +++ b/spec/lib/updater/git_updater_spec.rb @@ -44,4 +44,29 @@ describe GitUpdater do @git_updater.update().should === "Already up-to-date." end end + + describe "#has_local_changes?" do + after :each do + stub_system_command(@git_updater, /^git .* diff --exit-code 2>&1/, @stub_value) + @git_updater.has_local_changes?.should === @expected + end + + it "should return true if there are local changes" do + @stub_value = 'diff' + @expected = true + end + + it "should return false if there are no local changes" do + @stub_value = '' + @expected = false + end + end + + describe "#reset_head" do + it "should reset the local repo" do + stub_system_command(@git_updater, /^git .* reset --hard HEAD/, "HEAD is now at") + @git_updater.reset_head.should match(/^HEAD is now at/) + end + end + end diff --git a/wpscan.rb b/wpscan.rb index 1ed907d9..4a0bd989 100755 --- a/wpscan.rb +++ b/wpscan.rb @@ -52,6 +52,10 @@ begin # Check for updates if wpscan_options.update unless @updater.nil? + if @updater.has_local_changes? + puts "#{red('[!]')} Local file changes detected, an update will override local changes, do you want to continue updating? [y/n]" + Readline.readline =~ /^y/i ? @updater.reset_head : raise('Update aborted') + end puts @updater.update() else puts "Svn / Git not installed, or wpscan has not been installed with one of them."