From 2b85b44bd15c80ded3ac54051a40efb81f90ac80 Mon Sep 17 00:00:00 2001
From: g0tmi1k <have.you.g0tmi1k@gmail.com>
Date: Fri, 11 May 2018 11:19:51 +0100
Subject: [PATCH] Add offline database update support

---
 Gemfile                     |  1 +
 lib/common/common_helper.rb | 20 +++++++++++++-
 wpscan.rb                   | 52 +++++++++++++++++++++++++++----------
 3 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/Gemfile b/Gemfile
index fed1f734..b9001a26 100644
--- a/Gemfile
+++ b/Gemfile
@@ -6,6 +6,7 @@ gem 'addressable', '>=2.5.0'
 gem 'yajl-ruby', '>=1.3.0' # Better JSON parser regarding memory usage
 gem 'terminal-table', '>=1.6.0'
 gem 'ruby-progressbar', '>=1.8.1'
+gem 'rubyzip', '>=1.2.1'
 
 group :test do
   gem 'webmock', '>=2.3.2'
diff --git a/lib/common/common_helper.rb b/lib/common/common_helper.rb
index 05c6a96d..983ccec0 100644
--- a/lib/common/common_helper.rb
+++ b/lib/common/common_helper.rb
@@ -12,6 +12,7 @@ MODELS_LIB_DIR       = File.join(COMMON_LIB_DIR, 'models')
 COLLECTIONS_LIB_DIR  = File.join(COMMON_LIB_DIR, 'collections')
 
 DEFAULT_LOG_FILE     = File.join(ROOT_DIR, 'log.txt')
+DATA_FILE            = File.join(ROOT_DIR, 'data.zip') # wpscan/data.zip
 
 # Plugins directories
 COMMON_PLUGINS_DIR   = File.join(COMMON_LIB_DIR, 'plugins')
@@ -79,13 +80,30 @@ def add_trailing_slash(url)
   url =~ /\/$/ ? url : "#{url}/"
 end
 
-def missing_db_file?
+def missing_db_files?
   DbUpdater::FILES.each do |db_file|
     return true unless File.exist?(File.join(DATA_DIR, db_file))
   end
   false
 end
 
+# Find data.zip?
+def has_db_zip?
+  return File.exist?(DATA_FILE)? true : false
+end
+
+# Extract data.zip
+def extract_db_zip
+  puts DATA_FILE
+  Zip::File.open(DATA_FILE) do |zip_file|
+    zip_file.each do |f|
+      f_path = File.join(DATA_DIR, f.name)
+      FileUtils.mkdir_p(File.dirname(f_path))
+      zip_file.extract(f, f_path)
+    end
+  end
+end
+
 def last_update
   date = nil
   if File.exists?(LAST_UPDATE_FILE)
diff --git a/wpscan.rb b/wpscan.rb
index 09dc7b1e..988f954a 100755
--- a/wpscan.rb
+++ b/wpscan.rb
@@ -85,30 +85,56 @@ def main
       wpscan_options.to_h.merge(max_threads: wpscan_options.threads)
     )
 
-    # Check if db file needs upgrade (older than 5 days) and we are not running in --batch mode
+    # Check if database needs upgrade (if its older than 5 days) and we are not running in --batch mode
     # Also no need to check if the user supplied the --update switch
-    if update_required? && !wpscan_options.batch && !wpscan_options.update
+    if update_required? and  not wpscan_options.batch and not wpscan_options.update
+      # Banner
       puts
       puts notice('It seems like you have not updated the database for some time')
       puts notice("Last database update: #{date.strftime('%Y-%m-%d')}") unless date.nil?
+
+      # User prompt
       print '[?] Do you want to update now? [Y]es [N]o [A]bort, default: [N] > '
-      if (input = Readline.readline) =~ /^y/i
-        wpscan_options.update = true
-      elsif input =~ /^a/i
+      if (input = Readline.readline) =~ /^a/i
         puts 'Scan aborted'
         exit(1)
-      else
-        if missing_db_file?
-          puts critical('You can not run a scan without any databases. Manually extract the data.zip file.')
-          exit(1)
-        end
+      elsif input =~ /^y/i
+        wpscan_options.update = true
+      end
+
+      # Is there a database to go on with?
+      if missing_db_files? and not wpscan_options.update
+        puts critical('You can not run a scan without any databases')
+        exit(1)
       end
     end
 
+    # Should we update?
     if wpscan_options.update
-      puts notice('Updating the Database ...')
-      DbUpdater.new(DATA_DIR).update(wpscan_options.verbose)
-      puts notice('Update completed')
+      online_update = true
+
+      # Check for data.zip
+      if has_db_zip?
+        # User prompt
+        print '[?] Use the latest on-line database? Or use the off-line version? [O]n-line O[f]f-line [A]bort, default: [O] > '
+        if (input = Readline.readline) =~ /^a/i
+          puts 'Scan aborted'
+          exit(1)
+        elsif input =~ /^f/i
+          online_update = false
+        end
+      end
+
+      if online_update
+        puts notice('Updating the Database ...')
+        DbUpdater.new(DATA_DIR).update(wpscan_options.verbose)
+        puts notice('Update completed')
+      else
+        puts notice('Extracting the Database ...')
+        extract_db_zip
+        puts notice('Extraction completed')
+      end
+
       # Exit program if only option --update is used
       exit(0) unless wpscan_options.url
     end