diff --git a/db/migrations/005_add_logs.rb b/db/migrations/005_add_logs.rb new file mode 100644 index 0000000..095efdf --- /dev/null +++ b/db/migrations/005_add_logs.rb @@ -0,0 +1,16 @@ +# frozen_string_literal: true + +Sequel.migration do + up do + create_table :logs do + primary_key :id + + column :key, :string, size: 50, unique: true, null: false + column :value, :string, size: 100, null: false + end + end + + down do + drop_table :logs + end +end diff --git a/db/migrations/006_create_modules.rb b/db/migrations/006_create_modules.rb new file mode 100644 index 0000000..4226783 --- /dev/null +++ b/db/migrations/006_create_modules.rb @@ -0,0 +1,18 @@ +# frozen_string_literal: true + +Sequel.migration do + up do + create_table :modules do + primary_key :id + + column :path, :string, size: 255, null: false, unique: true + column :name, :string, size: 255, null: false + column :type, :string, size: 11, null: false + column :class_name, :string, size: 255, null: false, unique: true + end + end + + down do + drop_table :modules + end +end diff --git a/lib/cli/console.rb b/lib/cli/console.rb index 3015273..4d3b776 100644 --- a/lib/cli/console.rb +++ b/lib/cli/console.rb @@ -6,6 +6,7 @@ require 'cli/auto_complete' require 'cli/context' require 'cli/modules' +require 'cli/module_cache' require 'cli/module_info' require 'cli/loaded_module' require 'cli/options' @@ -19,6 +20,7 @@ class Console include Cli::AutoComplete include Cli::Help include Cli::Modules + include Cli::ModuleCache include Cli::LoadedModule include Cli::Options include Cli::Output @@ -50,10 +52,6 @@ def clear Gem.win_platform? ? (system 'cls') : (system 'clear') end - def quit - exit - end - def prompt_for_input prompt = 'wpxf'.underline.light_blue @@ -132,10 +130,22 @@ def execute_user_command(command, args) puts unless commands_without_output.include? command end + def check_cache + return if cache_valid? + + print_warning 'Refreshing the module cache...' + puts + rebuild_cache + end + def start + check_cache + loop do begin input = prompt_for_input + break if input =~ /exit|quit/ + command, *args = input.split(/\s/) execute_user_command(command, args) if command rescue StandardError => e diff --git a/lib/cli/module_cache.rb b/lib/cli/module_cache.rb new file mode 100644 index 0000000..12184a5 --- /dev/null +++ b/lib/cli/module_cache.rb @@ -0,0 +1,54 @@ +# frozen_string_literal: true + +module Cli + # A mixin to handle the database caching of module data. + module ModuleCache + def initialize + super + self.current_version_number = File.read(File.join(Wpxf.app_path, 'VERSION')) + end + + def cache_valid? + last_version_log = Models::Log.first(key: 'version') + return false if last_version_log.nil? + + current_version = Gem::Version.new(current_version_number) + last_version = Gem::Version.new(last_version_log.value) + + current_version == last_version + end + + def create_module_models(type) + namespace = type == 'exploit' ? Wpxf::Exploit : Wpxf::Auxiliary + namespace.module_list.each do |mod| + instance = mod[:class].new + + Models::Module.create( + path: mod[:name], + name: instance.module_name, + type: type, + class_name: mod[:class].to_s + ) + end + end + + def refresh_version_log + log = Models::Log.first(key: 'version') + log = Models::Log.new if log.nil? + log.key = 'version' + log.value = current_version_number + log.save + end + + def rebuild_cache + Models::Module.truncate + + create_module_models 'exploit' + create_module_models 'auxiliary' + + refresh_version_log + end + + attr_accessor :current_version_number + end +end diff --git a/lib/models/log.rb b/lib/models/log.rb new file mode 100644 index 0000000..fb3f376 --- /dev/null +++ b/lib/models/log.rb @@ -0,0 +1,20 @@ +# frozen_string_literal: true + +module Models + # A miscellaneous log entry. + class Log < Sequel::Model + plugin :validation_helpers + + def validate + super + + validates_presence :key + validates_type String, :key + validates_unique :key + validates_presence :value + + validates_max_length 50, :key + validates_max_length 100, :value + end + end +end diff --git a/lib/models/module.rb b/lib/models/module.rb new file mode 100644 index 0000000..f5a701a --- /dev/null +++ b/lib/models/module.rb @@ -0,0 +1,30 @@ +# frozen_string_literal: true + +module Models + # A cache of a {Wpxf::Module}'s metadata. + class Module < Sequel::Model + plugin :validation_helpers + + def validate + super + + validates_presence :path + validates_presence :name + validates_presence :type + validates_presence :class_name + + validates_type String, :path + validates_type String, :name + validates_type String, :class_name + + validates_unique :path + validates_unique :class_name + + validates_max_length 255, :path + validates_max_length 255, :name + validates_max_length 255, :class_name + + validates_format /^auxiliary|exploit$/, :type + end + end +end diff --git a/lib/wpxf/db.rb b/lib/wpxf/db.rb index 7246fad..b3768a2 100644 --- a/lib/wpxf/db.rb +++ b/lib/wpxf/db.rb @@ -9,4 +9,7 @@ module Db require 'models/workspace' require 'models/credential' +require 'models/log' +require 'models/module' + require 'wpxf/db/credentials' diff --git a/spec/lib/cli/console_spec.rb b/spec/lib/cli/console_spec.rb new file mode 100644 index 0000000..fda908b --- /dev/null +++ b/spec/lib/cli/console_spec.rb @@ -0,0 +1,56 @@ +# frozen_string_literal: true + +require_relative '../../spec_helper' +require 'cli/console' + +describe Cli::Console do + let(:subject) { Cli::Console.new } + + before :each, 'setup spies' do + allow(subject).to receive(:print_warning) + allow(subject).to receive(:rebuild_cache) + allow(subject).to receive(:puts) + end + + describe '#start' do + before :each, 'setup mocks' do + allow(subject).to receive(:prompt_for_input).and_return('exit') + allow(subject).to receive(:check_cache) + end + + it 'should check the module cache' do + subject.start + expect(subject).to have_received(:check_cache).exactly(1).times + end + end + + describe '#check_cache' do + context 'if the module cache is not valid' do + before(:each) do + allow(subject).to receive(:cache_valid?).and_return(false) + end + + it 'should refresh the module cache' do + subject.check_cache + expect(subject).to have_received(:rebuild_cache).exactly(1).times + end + + it 'should warn the user the cache is being refreshed' do + subject.check_cache + expect(subject).to have_received(:print_warning) + .with('Refreshing the module cache...') + end + end + + context 'if the module cache is valid' do + before(:each) do + allow(subject).to receive(:cache_valid?).and_return(true) + end + + it 'should not refresh the cache' do + subject.check_cache + expect(subject).to have_received(:rebuild_cache).exactly(0).times + end + end + end +end diff --git a/spec/lib/cli/module_cache_spec.rb b/spec/lib/cli/module_cache_spec.rb new file mode 100644 index 0000000..c7b3731 --- /dev/null +++ b/spec/lib/cli/module_cache_spec.rb @@ -0,0 +1,163 @@ +# frozen_string_literal: true + +require_relative '../../spec_helper' +require 'cli/module_cache' +require 'modules' + +describe Cli::ModuleCache do + let :subject do + Class.new do + include Cli::ModuleCache + + def initialize + super + end + end.new + end + + before(:each, 'setup mocks') do + allow(subject).to receive(:print_bad) + allow(subject).to receive(:print_good) + allow(subject).to receive(:print_warning) + allow(subject).to receive(:print_info) + allow(subject).to receive(:context) + end + + describe '#new' do + it 'should initialise `current_version_number` with the contents of the VERSION file' do + version_path = File.join(File.dirname(__FILE__), '../../../VERSION') + version = File.read(version_path) + expect(version).to match(/(\d\.?)+/) + expect(subject.current_version_number).to eql version + end + end + + describe '#cache_valid?' do + context 'if no version has been logged previously' do + it 'should return false' do + Models::Log.truncate + expect(subject.cache_valid?).to be false + end + end + + context 'if the previously logged version is older than the current' do + it 'should return false' do + Models::Log.create(key: 'version', value: '1.0') + subject.current_version_number = '1.1' + expect(subject.cache_valid?).to be false + end + end + + context 'if the current version is equal to the logged version' do + it 'should return true' do + Models::Log.create(key: 'version', value: '1.0') + subject.current_version_number = '1.0' + expect(subject.cache_valid?).to be true + end + end + + context 'if the current version is lower than the logged version' do + it 'should return false' do + Models::Log.create(key: 'version', value: '1.2') + subject.current_version_number = '1.1' + expect(subject.cache_valid?).to be false + end + end + end + + describe '#create_module_models' do + context 'if `type` is exploit' do + it 'should create a {Models::Module} for each class in the Wpxf::Exploit namespace' do + modules = Wpxf::Exploit.constants.select do |c| + Wpxf::Exploit.const_get(c).is_a? Class + end + + subject.create_module_models 'exploit' + exploit_count = Models::Module.where(type: 'exploit').count + expect(exploit_count).to eql modules.count + end + end + + context 'if `type` is not exploit' do + it 'should create a {Models::Module} for each class in the Wpxf::Exploit namespace' do + modules = Wpxf::Auxiliary.constants.select do |c| + Wpxf::Auxiliary.const_get(c).is_a? Class + end + + subject.create_module_models 'auxiliary' + exploit_count = Models::Module.where(type: 'auxiliary').count + expect(exploit_count).to eql modules.count + end + end + end + + describe '#refresh_version_log' do + context 'if a version log already exists' do + it 'should update the existing entry' do + Models::Log.create(key: 'version', value: '5') + expect(Models::Log.count).to eql 1 + + subject.current_version_number = '99' + subject.refresh_version_log + expect(Models::Log.count).to eql 1 + + log = Models::Log.first(key: 'version') + expect(log).to_not be_nil + expect(log.value.to_s).to eql '99' + end + end + + context 'if a version log does not exist' do + it 'should create a new entry' do + log = Models::Log.first(key: 'version') + expect(log).to be_nil + subject.current_version_number = '99' + subject.refresh_version_log + log = Models::Log.first(key: 'version') + expect(log).to_not be_nil + expect(log.value.to_s).to eql '99' + end + end + end + + describe '#rebuild_cache' do + it 'should truncate the existing cache' do + Models::Module.create( + path: 'exploit/shell/test', + name: 'test', + type: 'exploit', + class_name: 'Wpxf::Exploit::Test' + ) + + expect(Models::Module.first(name: 'test')).to_not be_nil + subject.rebuild_cache + expect(Models::Module.first(name: 'test')).to be_nil + end + + it 'should create a {Models::Module} for each class in the Wpxf::Exploit namespace' do + modules = Wpxf::Exploit.constants.select do |c| + Wpxf::Exploit.const_get(c).is_a? Class + end + + subject.rebuild_cache + exploit_count = Models::Module.where(type: 'exploit').count + expect(exploit_count).to eql modules.count + end + + it 'should create a {Models::Module} for each class in the Wpxf::Exploit namespace' do + modules = Wpxf::Auxiliary.constants.select do |c| + Wpxf::Auxiliary.const_get(c).is_a? Class + end + + subject.rebuild_cache + exploit_count = Models::Module.where(type: 'auxiliary').count + expect(exploit_count).to eql modules.count + end + + it 'should update the version log' do + allow(subject).to receive(:refresh_version_log) + subject.rebuild_cache + expect(subject).to have_received(:refresh_version_log).exactly(1).times + end + end +end diff --git a/spec/lib/models/log_spec.rb b/spec/lib/models/log_spec.rb new file mode 100644 index 0000000..3e6c1f0 --- /dev/null +++ b/spec/lib/models/log_spec.rb @@ -0,0 +1,11 @@ +# frozen_string_literal: true + +require_relative '../../spec_helper' + +describe Models::Log, type: :model do + it { is_expected.to validate_presence :key } + it { is_expected.to validate_presence :value } + it { is_expected.to validate_unique :key } + it { is_expected.to validate_max_length 50, :key } + it { is_expected.to validate_max_length 100, :value } +end diff --git a/spec/lib/models/module_spec.rb b/spec/lib/models/module_spec.rb new file mode 100644 index 0000000..5a1d623 --- /dev/null +++ b/spec/lib/models/module_spec.rb @@ -0,0 +1,19 @@ +# frozen_string_literal: true + +require_relative '../../spec_helper' + +describe Models::Module, type: :model do + it { is_expected.to validate_presence :path } + it { is_expected.to validate_presence :name } + it { is_expected.to validate_presence :type } + it { is_expected.to validate_presence :class_name } + + it { is_expected.to validate_unique :path } + it { is_expected.to validate_unique :class_name } + + it { is_expected.to validate_max_length 255, :path } + it { is_expected.to validate_max_length 255, :name } + it { is_expected.to validate_max_length 255, :class_name } + + it { is_expected.to validate_format(/^auxiliary|exploit$/, :type) } +end