Add --version and --update switches

Author: rastating

VERSION

@@ -0,0 +1 @@
+1.1.0

env.rb

@@ -20,6 +20,8 @@
 require 'wpxf/utility/text'
 require 'wpxf/utility/reference_inflater'
 
+require 'github_updater'
+
 module Wpxf
   def self.data_directory=(val)
     @@data_directory = val
@@ -29,6 +31,14 @@ def self.data_directory
     @@data_directory
   end
 
+  def self.app_path=(val)
+    @@app_path = val
+  end
+
+  def self.app_path
+    @@app_path
+  end
+
   def self.change_stdout_sync(enabled)
     original_setting = STDOUT.sync
     STDOUT.sync = true
@@ -37,4 +47,5 @@ def self.change_stdout_sync(enabled)
   end
 end
 
+Wpxf.app_path = app_path
 Wpxf.data_directory = File.join(app_path, 'data/')

github_updater.rb

@@ -0,0 +1,81 @@
+require 'json'
+require 'typhoeus'
+require 'fileutils'
+
+module Wpxf
+  # A self updater that uses the latest release from GitHub as the target.
+  class GitHubUpdater
+    def latest_release_url
+      'https://api.github.com/repos/rastating/wordpress-exploit-framework/releases/latest'
+    end
+
+    # Get information about the latest update available on GitHub.
+    # @param current_version [String] the current version number in use.
+    # @return [Hash, nil] a hash containing the :release_notes, :zip_url and :release_name, or nil if there is no update.
+    def get_update(current_version)
+      res = Typhoeus.get(latest_release_url)
+      return nil unless res && res.code == 200
+
+      begin
+        update = JSON.parse(res.body)
+      rescue JSON::ParserError
+        return nil
+      end
+
+      begin
+        return nil if Gem::Version.new(current_version) >= Gem::Version.new(update['tag_name'].sub(/^v/, ''))
+      rescue
+        return nil
+      end
+
+      { release_notes: update['body'], zip_url: update['zipball_url'], release_name: update['name'] }
+    end
+
+    # Download and apply an update from the specified URL.
+    # @param zip_url [String] the URL to fetch the update from.
+    def download_and_apply_update(zip_url)
+      tmp = create_tmp_directory
+      zip_filename = File.join(tmp, 'update.zip')
+
+      download_update_zip(zip_url, zip_filename)
+
+      Zip::File.open(zip_filename) do |zip_file|
+        zip_file.each do |entry|
+          entry.extract File.join(tmp, entry.name)
+        end
+      end
+
+      Dir.glob(File.join(tmp, 'rastating-wordpress-exploit-framework*/*')) do |f|
+        FileUtils.cp_r(f, Wpxf.app_path)
+      end
+
+      FileUtils.rm_rf(tmp)
+    end
+
+    private
+
+    def create_tmp_directory
+      tmp = File.join(Dir.tmpdir, "wpxf_update_#{object_id}")
+      FileUtils.mkdir_p(tmp)
+      tmp
+    end
+
+    def download_update_zip(zip_url, local_filename)
+      zip = File.open(local_filename, 'wb')
+      request = Typhoeus::Request.new(zip_url, followlocation: true)
+      request.on_headers do |response|
+        raise 'Request failed' if response.code != 200
+      end
+
+      request.on_body do |chunk|
+        zip.write(chunk)
+      end
+
+      request.on_complete do
+        zip.close
+      end
+
+      request.run
+    end
+  end
+end

spec/github_updater_spec.rb

@@ -0,0 +1,87 @@
+require_relative 'spec_helper'
+
+describe Wpxf::GitHubUpdater do
+  let(:typhoeus_return_code) { :ok }
+  let(:typhoeus_code) { 200 }
+  let(:typhoeus_body) { '' }
+  let(:typhoeus_headers) { { 'Content-Type' => 'text/html; charset=utf-8' } }
+  let(:subject) { Wpxf::GitHubUpdater.new }
+
+  before :each do
+    Typhoeus.stub(/.*/) do
+      Typhoeus::Response.new(
+        code: typhoeus_code,
+        body: typhoeus_body,
+        headers: typhoeus_headers,
+        return_code: typhoeus_return_code
+      )
+    end
+  end
+
+  describe '#get_update' do
+    context 'when the API returns a status other than 200' do
+      let(:typhoeus_code) { 404 }
+
+      it 'returns nil' do
+        expect(subject.get_update('')).to be_nil
+      end
+    end
+
+    context 'when invalid JSON is returned from the API' do
+      let(:typhoeus_body) { 'invalid body' }
+
+      it 'returns nil' do
+        expect(subject.get_update('')).to be_nil
+      end
+    end
+
+    context 'when the version number from GitHub is invalid' do
+      let(:typhoeus_body) do
+        { tag_name: 'invalid version' }.to_json
+      end
+
+      it 'returns nil' do
+        expect(subject.get_update('1.0')).to be_nil
+      end
+    end
+
+    context 'when the current version number is invalid' do
+      let(:typhoeus_body) do
+        { tag_name: '1.1' }.to_json
+      end
+
+      it 'returns nil' do
+        expect(subject.get_update('invalid')).to be_nil
+      end
+    end
+
+    context 'when the current version is the same as the latest GitHub release' do
+      let(:typhoeus_body) do
+        { tag_name: '1.1' }.to_json
+      end
+
+      it 'returns nil' do
+        expect(subject.get_update('1.1')).to be_nil
+      end
+    end
+
+    context 'when the current version is older than the latest GitHub release' do
+      let(:typhoeus_body) do
+        {
+          tag_name: 'v1.1',
+          body: 'notes',
+          zipball_url: 'url',
+          name: 'v1.1'
+        }.to_json
+      end
+
+      it 'returns a hash containing the information about the latest update' do
+        expect(subject.get_update('1.0')).to include(
+          release_notes: 'notes',
+          zip_url: 'url',
+          release_name: 'v1.1'
+        )
+      end
+    end
+  end
+end

wpxf.rb

@@ -3,6 +3,38 @@
 require 'fileutils'
 require_relative 'env'
 require 'cli/console'
+require 'slop'
+
+Slop.parse do |o|
+  o.on '--update', 'check for updates' do
+    current_version = File.read('VERSION').strip
+
+    updater = Wpxf::GitHubUpdater.new
+    update = updater.get_update(current_version)
+
+    if update.nil?
+      puts 'No updates available'
+      exit
+    end
+
+    puts 'A new update is available!'
+    puts
+    puts '-- Release Notes --'
+    puts update[:release_notes]
+    puts
+    puts "Downloading latest update (#{update[:release_name]})..."
+    updater.download_and_apply_update(update[:zip_url])
+
+    puts 'Update finished! Make sure to run "bundle install" in the WPXF directory.'
+    puts
+    exit
+  end
+
+  o.on '--version', 'print the version' do
+    puts File.read('VERSION').strip
+    exit
+  end
+end
 
 puts '                           _'
 puts '    __      _____  _ __ __| |_ __  _ __ ___  ___ ___'
@@ -31,7 +63,7 @@
 
 Dir.chdir(Dir.tmpdir) do
   temp_directories = Dir.glob('wpxf_*')
-  if temp_directories.length > 0
+  unless temp_directories.empty?
     print '[!] '.yellow
     puts "#{temp_directories.length} temporary files were found that "\
          'appear to no longer be needed.'