Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions spec/beef/core/main/autorun_engine/autorun_engine_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion spec/beef/core/main/handlers/browser_details_handler_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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.'
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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'
Expand All @@ -25,6 +26,7 @@

after(:all) do
Process.kill("INT",@pid)
SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
end

it 'delete' do
Expand Down
2 changes: 2 additions & 0 deletions spec/beef/core/main/network_stack/handlers/redirector_spec.rb
Original file line number Diff line number Diff line change
@@ -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'
Expand All @@ -26,6 +27,7 @@

after(:all) do
Process.kill("INT",@pid)
SpecActiveRecordConnection.restore!(@__ar_config_snapshot)
end

it 'redirects' do
Expand Down
2 changes: 2 additions & 0 deletions spec/beef/extensions/requester_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down Expand Up @@ -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
2 changes: 2 additions & 0 deletions spec/beef/extensions/websocket_hooked_browser_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions spec/beef/modules/debug/test_beef_debugs_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
58 changes: 58 additions & 0 deletions spec/spec_helper.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading