Fix broken namespaces and migrate entry point into new bin

rastating · Jul 14, 2018 · 3fe58dd · 3fe58dd
1 parent c0f1f98
commit 3fe58dd
Show file tree

Hide file tree

Showing 42 changed files with 219 additions and 297 deletions.
diff --git a/bin/wpxf b/bin/wpxf
@@ -1 +1,52 @@
 #!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require 'wpxf'
+require 'wpxf/cli/console'
+require 'wpxf/cli/banner'
+
+begin
+  Slop.parse do |o|
+    o.on '--version', 'print the version' do
+      puts Wpxf.version
+      exit
+    end
+  end
+rescue Slop::UnknownOption => e
+  puts e.message
+  exit
+end
+
+console = Wpxf::Cli::Console.new
+console.check_cache
+console.clear
+
+banner = Wpxf::Cli::Banner.new
+banner.display
+
+Dir.chdir(Dir.tmpdir) do
+  temp_directories = Dir.glob('wpxf_*')
+  unless temp_directories.empty?
+    print '[!] '.yellow
+    puts "#{temp_directories.length} temporary files were found that "\
+         'appear to no longer be needed.'
+    print '    Would you like to remove these files? [y/n]: '
+    temp_directories.each { |d| FileUtils.rm_r(d) } if gets.chomp =~ /^y$/i
+    puts
+  end
+end
+
+found_env_var = false
+ENV.each do |name, value|
+  next if name.casecmp('wpxf_env').zero?
+  match = name.match(/^wpxf_(.+)/i)
+
+  if match
+    console.gset match.captures[0], value
+    found_env_var = true
+  end
+end
+
+puts if found_env_var
+console.start
+console.clear
diff --git a/env.rb b/env.rb
diff --git a/lib/wpxf.rb b/lib/wpxf.rb
@@ -49,3 +49,4 @@ def self.change_stdout_sync(enabled)
 end
 
 require_relative '../db/env'
+require 'wpxf/core'
diff --git a/lib/wpxf/cli/auto_complete.rb b/lib/wpxf/cli/auto_complete.rb
@@ -51,7 +51,7 @@ def refresh_autocomplete_options
       def build_cmd_list
         cmds = {}
         permitted_commands.each { |c| cmds[c] = {} }
-        Models::Module.each { |m| cmds['use'][m.path] = {} }
+        Wpxf::Models::Module.each { |m| cmds['use'][m.path] = {} }
         cmds['show'] = {
           'options' => {},
           'advanced' => {},

diff --git a/lib/wpxf/cli/banner.rb b/lib/wpxf/cli/banner.rb
@@ -20,11 +20,11 @@ def format_colour(value)
       end
 
       def auxiliary_count
-        Models::Module.where(type: 'auxiliary').count
+        Wpxf::Models::Module.where(type: 'auxiliary').count
       end
 
       def exploit_count
-        Models::Module.where(type: 'exploit').count
+        Wpxf::Models::Module.where(type: 'exploit').count
       end
 
       def format_data(value)

diff --git a/lib/wpxf/cli/context.rb b/lib/wpxf/cli/context.rb
@@ -18,7 +18,7 @@ def load_module(path)
       end
 
       def reload
-        load("#{@module_path}.rb")
+        load("wpxf/modules/#{@module_path}.rb")
         load_module(@module_path)
       end
 

diff --git a/lib/wpxf/cli/creds.rb b/lib/wpxf/cli/creds.rb
@@ -5,7 +5,7 @@ module Cli
     # A mixin to provide functions to interact with credentials stored in the database.
     module Creds
       def delete_credential(id)
-        item = Models::Credential.first(
+        item = Wpxf::Models::Credential.first(
           id: id.to_i,
           workspace: active_workspace
         )
@@ -28,7 +28,7 @@ def list_credentials
           type: 'Type'
         )
 
-        Models::Credential.where(workspace: active_workspace).each do |cred|
+        Wpxf::Models::Credential.where(workspace: active_workspace).each do |cred|
           rows.push(
             id: cred.id,
             host: "#{cred.host}:#{cred.port}",

diff --git a/lib/wpxf/cli/help.rb b/lib/wpxf/cli/help.rb
@@ -65,7 +65,7 @@ def help
       end
 
       def show_exploits
-        modules = Models::Module.where(type: 'exploit')
+        modules = Wpxf::Models::Module.where(type: 'exploit')
                                 .order(:path)
                                 .map { |m| { path: m.path, title: m.name } }
 
@@ -74,7 +74,7 @@ def show_exploits
       end
 
       def show_auxiliary
-        modules = Models::Module.where(type: 'auxiliary')
+        modules = Wpxf::Models::Module.where(type: 'auxiliary')
                                 .order(:path)
                                 .map { |m| { path: m.path, title: m.name } }
 

diff --git a/lib/wpxf/cli/loot.rb b/lib/wpxf/cli/loot.rb
@@ -18,7 +18,7 @@ def build_loot_table
           type: 'Type'
         )
 
-        Models::LootItem.where(workspace: active_workspace).each do |item|
+        Wpxf::Models::LootItem.where(workspace: active_workspace).each do |item|
           rows.push(
             id: item.id,
             host: "#{item.host}:#{item.port}",
@@ -39,7 +39,7 @@ def list_loot
       end
 
       def find_loot_item(id)
-        item = Models::LootItem.first(
+        item = Wpxf::Models::LootItem.first(
           id: id.to_i,
           workspace: active_workspace
         )

diff --git a/lib/wpxf/cli/module_cache.rb b/lib/wpxf/cli/module_cache.rb
@@ -6,11 +6,11 @@ module Cli
     module ModuleCache
       def initialize
         super
-        self.current_version_number = File.read(File.join(Wpxf.app_path, 'VERSION'))
+        self.current_version_number = Wpxf.version
       end
 
       def cache_valid?
-        last_version_log = Models::Log.first(key: 'version')
+        last_version_log = Wpxf::Models::Log.first(key: 'version')
         return false if last_version_log.nil?
 
         current_version = Gem::Version.new(current_version_number)
@@ -24,7 +24,7 @@ def create_module_models(type)
         namespace.module_list.each do |mod|
           instance = mod[:class].new
 
-          Models::Module.create(
+          Wpxf::Models::Module.create(
             path: mod[:name],
             name: instance.module_name,
             type: type,
@@ -34,8 +34,8 @@ def create_module_models(type)
       end
 
       def refresh_version_log
-        log = Models::Log.first(key: 'version')
-        log = Models::Log.new if log.nil?
+        log = Wpxf::Models::Log.first(key: 'version')
+        log = Wpxf::Models::Log.new if log.nil?
         log.key = 'version'
         log.value = current_version_number
         log.save
@@ -44,7 +44,7 @@ def refresh_version_log
       def rebuild_cache
         print_warning 'Refreshing the module cache...'
 
-        Models::Module.truncate
+        Wpxf::Models::Module.truncate
 
         create_module_models 'exploit'
         create_module_models 'auxiliary'

diff --git a/lib/wpxf/cli/modules.rb b/lib/wpxf/cli/modules.rb
@@ -59,7 +59,7 @@ def module_name_from_class(klass)
       end
 
       def search_modules(args)
-        modules = Models::Module.where(Sequel.ilike(:name, "%#{args.join(' ')}%"))
+        modules = Wpxf::Models::Module.where(Sequel.ilike(:name, "%#{args.join(' ')}%"))
         modules.map { |m| { path: m.path, title: m.name } }
       end
 

diff --git a/lib/wpxf/cli/workspace.rb b/lib/wpxf/cli/workspace.rb
@@ -7,7 +7,7 @@ module Workspace
       def initialize
         super
 
-        self.active_workspace = Models::Workspace.first(name: 'default')
+        self.active_workspace = Wpxf::Models::Workspace.first(name: 'default')
       end
 
       def workspace(*args)
@@ -24,16 +24,16 @@ def workspace(*args)
       end
 
       def workspaces
-        Models::Workspace.all
+        Wpxf::Models::Workspace.all
       end
 
       def add_workspace(name)
-        unless Models::Workspace.where(name: name).count.zero?
+        unless Wpxf::Models::Workspace.where(name: name).count.zero?
           return print_warning "#{name} already exists"
         end
 
         begin
-          Models::Workspace.create(name: name)
+          Wpxf::Models::Workspace.create(name: name)
           return print_good "Added workspace: #{name}"
         rescue Sequel::ValidationFailed
           print_warning 'Workspace names may only contain 1-50 alphanumeric characters and underscores'
@@ -51,7 +51,7 @@ def list_workspaces
       end
 
       def switch_workspace(name)
-        next_workspace = Models::Workspace.first(name: name)
+        next_workspace = Wpxf::Models::Workspace.first(name: name)
 
         if next_workspace
           self.active_workspace = next_workspace
@@ -69,7 +69,7 @@ def delete_workspace(name)
         end
 
         current_name = active_workspace.name
-        Models::Workspace.where(name: name).destroy
+        Wpxf::Models::Workspace.where(name: name).destroy
         print_good "Deleted workspace: #{name}"
         switch_workspace 'default' if name == current_name
       end

diff --git a/lib/wpxf/db/credentials.rb b/lib/wpxf/db/credentials.rb
@@ -8,15 +8,15 @@ module Wpxf::Db::Credentials
   # @param type [String] the type of string stored in the password field.
   # @return [Models::Credential] the newly created {Models::Credential}.
   def store_credentials(username, password = '', type = 'plain')
-    credential = Models::Credential.first(
+    credential = Wpxf::Models::Credential.first(
       host: target_host,
       port: target_port,
       username: username,
       type: type,
       workspace: active_workspace
     )
 
-    credential = Models::Credential.new if credential.nil?
+    credential = Wpxf::Models::Credential.new if credential.nil?
     credential.host = target_host
     credential.port = target_port
     credential.username = username

diff --git a/lib/wpxf/db/loot.rb b/lib/wpxf/db/loot.rb
@@ -7,7 +7,7 @@ module Wpxf::Db::Loot
   # @param type [String] the type of loot acquired.
   # @return [Models::LootItem] the newly created {Models::LootItem}.
   def store_loot(path, notes = '', type = 'unknown')
-    Models::LootItem.create(
+    Wpxf::Models::LootItem.create(
       host: target_host,
       port: target_port,
       path: path,

diff --git a/lib/wpxf/modules.rb b/lib/wpxf/modules.rb
@@ -6,7 +6,7 @@ def self.build_module_list(namespace, folder_name = '')
       namespace.const_get(c).is_a? Class
     end
 
-    modules_directory = File.join(Wpxf.app_path, folder_name)
+    modules_directory = File.join(Wpxf.app_path, 'lib', 'wpxf', folder_name)
 
     modules.map do |m|
       klass = namespace.const_get(m)
@@ -19,7 +19,7 @@ def self.build_module_list(namespace, folder_name = '')
   end
 
   def self.load_module(path)
-    mod = Models::Module.first(path: path)
+    mod = Wpxf::Models::Module.first(path: path)
     raise "\"#{path}\" is not a valid module" if mod.nil?
     Object.const_get(mod.class_name).new
   end

diff --git a/spec/lib/wpxf/cli/banner_spec.rb b/spec/lib/wpxf/cli/banner_spec.rb
@@ -10,7 +10,7 @@
     allow(subject).to receive(:puts)
 
     (1..5).each do |i|
-      Models::Module.create(
+      Wpxf::Models::Module.create(
         path: "test#{i}/",
         name: "Test #{i}",
         class_name: "Test#{i}",
@@ -19,7 +19,7 @@
     end
 
     (6..8).each do |i|
-      Models::Module.create(
+      Wpxf::Models::Module.create(
         path: "test#{i}/",
         name: "Test #{i}",
         class_name: "Test#{i}",
@@ -30,7 +30,7 @@
 
   describe '#new' do
     it 'should initialise #raw_content with the content of `data/banners/default.txt`' do
-      base_path = File.expand_path(File.join(__dir__, '..', '..', '..'))
+      base_path = File.expand_path(File.join(__dir__, '..', '..', '..', '..'))
       banner_path = File.join(base_path, 'data', 'banners', 'default.txt')
       expected = File.read(banner_path)
 
@@ -79,10 +79,7 @@
 
   describe '#format_data' do
     it 'should replace {VERSION} with the current version of WPXF' do
-      base_path = File.expand_path(File.join(__dir__, '..', '..', '..'))
-      version_path = File.join(base_path, 'VERSION')
-      expected = File.read(version_path).strip
-
+      expected = Wpxf.version
       expect(subject.format_data('V:{VERSION}')).to eql "V:#{expected}"
     end