From 3e56a945730accc2999d6a7e7087936ea70a5c49 Mon Sep 17 00:00:00 2001
From: Stephen <stephenahq@gmail.com>
Date: Mon, 25 Mar 2024 07:26:54 +1000
Subject: [PATCH] tidy up auth_rate_spec

---
 spec/beef/api/auth_rate_spec.rb | 19 +--------------
 spec/spec_helper.rb             | 42 ++++++++++++++++-----------------
 2 files changed, 22 insertions(+), 39 deletions(-)

diff --git a/spec/beef/api/auth_rate_spec.rb b/spec/beef/api/auth_rate_spec.rb
index 06f36a745..4a7589882 100644
--- a/spec/beef/api/auth_rate_spec.rb
+++ b/spec/beef/api/auth_rate_spec.rb
@@ -4,22 +4,15 @@
 # See the file 'doc/COPYING' for copying permission
 #
 
-require 'net/http'
-require 'uri'
-
 RSpec.describe 'BeEF API Rate Limit' do
-    
+
 	before(:each) do
-        port = 3000
-        expect(`lsof -i :#{port}`).to be_empty
         @pid = start_beef_server_and_wait
         @username = @config.get('beef.credentials.user')
         @password = @config.get('beef.credentials.passwd')
 	end
 
 	after(:each) do
-        # Stop the DNS server after each test case
-        BeEF::Extension::Dns::Server.instance.stop
 		# Shutting down server
         Process.kill("KILL", @pid) unless @pid.nil?
         Process.wait(@pid) unless @pid.nil? # Ensure the process has exited and the port is released 
@@ -27,19 +20,11 @@ RSpec.describe 'BeEF API Rate Limit' do
 	end
 
     it 'confirm correct creds are successful' do
-
-        # sleep 1
-        # uri = URI.parse("http://#{ATTACK_DOMAIN}:3000")
-        # response = Net::HTTP.get_response(uri)
-        # expect(response).to be_a(Net::HTTPSuccess) # HTTP request is successful
-
-        # sleep 60
         test_api = BeefRestClient.new('http', ATTACK_DOMAIN, '3000', @username, @password) 
         expect(@config.get('beef.credentials.user')).to eq('beef')
         expect(@config.get('beef.credentials.passwd')).to eq('beef')
         expect(test_api.auth()[:payload]).not_to eql("401 Unauthorized") 
         expect(test_api.auth()[:payload]["success"]).to be(true) # valid pass should succeed
-
     end
     
     it 'confirm incorrect creds are unsuccessful' do
@@ -49,7 +34,6 @@ RSpec.describe 'BeEF API Rate Limit' do
     end
     
     it 'adheres to 9 bad passwords then 1 correct auth rate limits' do
-
         # create api structures with bad passwords and one good
 		passwds = (1..9).map { |i| "bad_password"} # incorrect password
 		passwds.push @password # correct password
@@ -62,7 +46,6 @@ RSpec.describe 'BeEF API Rate Limit' do
     end
     
     it 'adheres to random bad passords and 1 correct auth rate limits' do
-
         # create api structures with bad passwords and one good
 		passwds = (1..9).map { |i| "bad_password"} # incorrect password
 		passwds.push @password # correct password
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 320587ed8..7c0a240fd 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -90,14 +90,12 @@ RSpec.configure do |config|
     Process.kill('KILL', server_pids)
   end
 
-
-
 ########################################
 
 require 'socket'
 
   def port_available?
-    socket = TCPSocket.new('localhost', 3000)
+    socket = TCPSocket.new(@host, @port)
     socket.close
     false  # If a connection is made, the port is in use, so it's not available.
   rescue Errno::ECONNREFUSED
@@ -106,17 +104,13 @@ require 'socket'
     true   # If the connection is refused, the port is not in use, so it's available.
   end
 
-
   def configure_beef
-
     # Reset or re-initialise the configuration to a default state
     @config = BeEF::Core::Configuration.instance
 
     @config.set('beef.credentials.user', "beef")
     @config.set('beef.credentials.passwd', "beef")
-
-    @username = @config.get('beef.credentials.user')
-    @password = @config.get('beef.credentials.passwd')
+    @config.set('beef.http.https.enable', false)
   end
 
   # Load the server
@@ -129,10 +123,14 @@ require 'socket'
   end
 
   def start_beef_server
-    exit unless port_available?
     configure_beef
-    load_beef_extensions_and_modules
+    @port = @config.get('beef.http.port')
+    @host = @config.get('beef.http.host')
+    @host = '127.0.0.1'
 
+    exit unless port_available?
+    load_beef_extensions_and_modules
+    
     # Grab DB file and regenerate if requested
     db_file = @config.get('beef.database.file')
 
@@ -176,27 +174,29 @@ require 'socket'
   end
 
   def beef_server_running?(uri_str)
-    uri = URI.parse(uri_str)
-    response = Net::HTTP.get_response(uri)
-    response.is_a?(Net::HTTPSuccess)
-  rescue
-    false
+    begin
+      uri = URI.parse(uri_str)
+      response = Net::HTTP.get_response(uri)
+      response.is_a?(Net::HTTPSuccess)
+      rescue Errno::ECONNREFUSED
+        return false # Connection refused means the server is not running
+      rescue StandardError => e
+        return false # Any other error means the server is not running
+    end
   end
 
   def wait_for_beef_server_to_start(uri_str, timeout: 5)
-    start_time = Time.now
-
+    start_time = Time.now # Record the time we started
     until beef_server_running?(uri_str) || (Time.now - start_time) > timeout do
-      sleep 0.1
+      sleep 0.1 # Wait a bit before checking again
     end
-
-    beef_server_running?(uri_str)
+    beef_server_running?(uri_str) # Return the result of the check
   end
 
   def start_beef_server_and_wait
     pid = start_beef_server
 
-    if wait_for_beef_server_to_start('http://localhost:3000', timeout: 5)
+    if wait_for_beef_server_to_start('http://localhost:3000', timeout: 3)
       # print_info "Server started successfully."
     else
       print_error "Server failed to start within timeout."