Merge pull request #3462 from kaitozaw/test/rspec-isolation

Fixes #3464

Author: zinduolis

spec/beef/core/main/autorun_engine/autorun_engine_spec.rb

@@ -12,6 +12,7 @@
 
 RSpec.describe 'AutoRunEngine Test', run_on_browserstack: true do
   before(:all) do
+    @__ar_config_snapshot = SpecActiveRecordConnection.snapshot
     @config = BeEF::Core::Configuration.instance
 
     # Grab DB file and regenerate if requested
@@ -110,6 +111,7 @@
   after(:all) do
     server_teardown(@driver, @pid, @pids)
     disconnect_all_active_record!
+    SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
   end
 
   it 'AutoRunEngine is working' do

spec/beef/core/main/handlers/browser_details_handler_spec.rb

@@ -12,7 +12,7 @@
 
 RSpec.describe 'Browser Details Handler', run_on_browserstack: true do
   before(:all) do
-
+    @__ar_config_snapshot = SpecActiveRecordConnection.snapshot
     @config = BeEF::Core::Configuration.instance
     db_file = @config.get('beef.database.file')
     print_info 'Resetting the database for BeEF.'
@@ -106,6 +106,7 @@
   after(:all) do
     server_teardown(@driver, @pid, @pids)
     disconnect_all_active_record!
+    SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
   end
 
   it 'can successfully hook a browser' do

spec/beef/core/main/network_stack/handlers/dynamic_reconstruction_spec.rb

@@ -1,6 +1,7 @@
 RSpec.describe 'BeEF Dynamic Reconsturction' do
 
   before(:all) do
+    @__ar_config_snapshot = SpecActiveRecordConnection.snapshot
     @port = 2001
     config = {}
     config[:BindAddress] = '127.0.0.1'
@@ -25,6 +26,7 @@
 
   after(:all) do
     Process.kill("INT",@pid)
+    SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
   end
 
   it 'delete' do

spec/beef/core/main/network_stack/handlers/redirector_spec.rb

@@ -1,6 +1,7 @@
 RSpec.describe 'BeEF Redirector' do
 
   before(:all) do
+    @__ar_config_snapshot = SpecActiveRecordConnection.snapshot
     @port = 2002
     config = {}
     config[:BindAddress] = '127.0.0.1'
@@ -26,6 +27,7 @@
 
   after(:all) do
     Process.kill("INT",@pid)
+    SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
   end
 
   it 'redirects' do

spec/beef/extensions/requester_spec.rb

@@ -24,6 +24,7 @@
 
   xit 'requester works' do
     begin
+      ar_snapshot = SpecActiveRecordConnection.snapshot
       # Start beef server
       @config = BeEF::Core::Configuration.instance
       @config.set('beef.credentials.user', 'beef')
@@ -77,6 +78,7 @@
       BeEF::Core::Models::Http.where(hooked_browser_id: hb_session).delete_all if defined? hb_session
       Process.kill('KILL', @pid) if defined? @pid
       Process.kill('KILL', @pids) if defined? @pids
+      SpecActiveRecordConnection.restore!(ar_snapshot)
     end
   end
 end

spec/beef/extensions/websocket_hooked_browser_spec.rb

@@ -13,6 +13,7 @@
 
 RSpec.describe 'Browser hooking with Websockets', run_on_browserstack: true do
   before(:all) do
+    @__ar_config_snapshot = SpecActiveRecordConnection.snapshot
     @config = BeEF::Core::Configuration.instance
     # Grab DB file and regenerate if requested
     print_info 'Loading database'
@@ -104,6 +105,7 @@
   after(:all) do
     server_teardown(@driver, @pid, @pids)
     disconnect_all_active_record!
+    SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
   end
 
   it 'confirms a websocket server has been started' do

spec/beef/modules/debug/test_beef_debugs_spec.rb

@@ -12,6 +12,7 @@
 
 RSpec.describe 'BeEF Debug Command Modules:', run_on_browserstack: true do
   before(:all) do
+    @__ar_config_snapshot = SpecActiveRecordConnection.snapshot
     # Grab config and set creds in variables for ease of access
     @config = BeEF::Core::Configuration.instance
     @pids = []  # ensure defined for teardown consistency
@@ -128,6 +129,7 @@
   after(:all) do
     server_teardown(@driver, @pid, @pids)
     disconnect_all_active_record!
+    SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
   end
 
   it 'The Test_beef.debug() command module successfully executes' do

spec/spec_helper.rb

@@ -271,6 +271,8 @@ def start_beef_server
     # Generate a token for the server to respond with
     BeEF::Core::Crypto::api_token
 
+    disconnect_all_active_record!
+
     # Initiate server start-up
     BeEF::API::Registrar.instance.fire(BeEF::API::Server, 'pre_http_start', http_hook_server)
     pid = fork do
@@ -322,4 +324,60 @@ def stop_beef_server(pid)
     pid = nil       
   end
 
+end
+
+# -------------------------------------------------------------------
+# ActiveRecord connection snapshot/restore helpers (test isolation)
+# Some specs disconnect ActiveRecord (fork safety), destroying the SQLite in-memory DB.
+# These helpers restore it for later specs.
+# -------------------------------------------------------------------
+module SpecActiveRecordConnection
+  module_function
+
+  def snapshot
+    # Capture the current AR connection configuration hash if possible.
+    if ActiveRecord::Base.respond_to?(:connection_db_config) && ActiveRecord::Base.connection_db_config
+      ActiveRecord::Base.connection_db_config.configuration_hash
+    else
+      ActiveRecord::Base.connection_config
+    end
+  rescue StandardError
+    nil
+  end
+
+  def restore!(config_hash)
+    # Ensure we don't leave AR disconnected for subsequent specs.
+    begin
+      handler = ActiveRecord::Base.connection_handler
+      if handler.respond_to?(:connection_pool_list)
+        handler.connection_pool_list.each { |pool| pool.disconnect! }
+      elsif handler.respond_to?(:connection_pools)
+        handler.connection_pools.each_value { |pool| pool.disconnect! }
+      else
+        ActiveRecord::Base.connection_pool.disconnect!
+      end
+    rescue StandardError
+      # ignore
+    end
+
+    if config_hash
+      OTR::ActiveRecord.configure_from_hash!(config_hash)
+    else
+      # Fallback to suite default
+      OTR::ActiveRecord.configure_from_hash!(adapter: 'sqlite3', database: ':memory:')
+    end
+
+    if Gem.loaded_specs['otr-activerecord'].version > Gem::Version.create('1.4.2')
+      OTR::ActiveRecord.establish_connection!
+    end
+    ActiveRecord::Schema.verbose = false
+
+    # Run migrations if the restored DB is empty/outdated
+    ActiveRecord::Migration.verbose = false
+    ActiveRecord::Migrator.migrations_paths = [File.join('core', 'main', 'ar-migrations')]
+    context = ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths)
+    if context.needs_migration?
+      ActiveRecord::Migrator.new(:up, context.migrations, context.schema_migration, context.internal_metadata).migrate
+    end
+  end
 end