From 065c7adf036c8b02c52fbe1927667068621bec30 Mon Sep 17 00:00:00 2001
From: kaitoozawa <kaitozaw@gmail.com>
Date: Sat, 3 Jan 2026 15:56:46 +1000
Subject: [PATCH] fix test cleanup and install firefox dependencies

---
 Dockerfile                 |  8 ++++++++
 spec/spec_helper.rb        |  7 ++-----
 spec/support/ui_support.rb | 35 +++++++++++++++++++++--------------
 3 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 1207287a5..e70bcccc2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -78,9 +78,17 @@ RUN adduser --home /beef --gecos beef --disabled-password beef \
     zlib1g \
     bison \
     nodejs \
+    firefox-esr \
  && apt-get -y clean \
  && rm -rf /var/lib/apt/lists/*
 
+# Install geckodriver for Selenium tests
+RUN GECKODRIVER_VERSION=$(curl -sL https://api.github.com/repos/mozilla/geckodriver/releases/latest | grep -oP '"tag_name": "\K(.*)(?=")') \
+ && wget -q "https://github.com/mozilla/geckodriver/releases/download/${GECKODRIVER_VERSION}/geckodriver-${GECKODRIVER_VERSION}-linux64.tar.gz" \
+ && tar -xzf "geckodriver-${GECKODRIVER_VERSION}-linux64.tar.gz" -C /usr/local/bin \
+ && chmod +x /usr/local/bin/geckodriver \
+ && rm "geckodriver-${GECKODRIVER_VERSION}-linux64.tar.gz"
+
 # Use gemset created by the builder above
 COPY --chown=beef:beef . /beef
 COPY --from=builder /usr/local/bundle /usr/local/bundle
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 8bdc9b7d7..ee4416be9 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -229,8 +229,7 @@ require 'socket'
     @host = '127.0.0.1'
 
     unless port_available?
-      print_error "Port #{@port} is already in use. Exiting."
-      exit
+      raise "Port #{@port} is already in use. Cannot start BeEF server."
     end
     load_beef_extensions_and_modules
     
@@ -317,11 +316,9 @@ require 'socket'
   end
 
   def stop_beef_server(pid)
-    exit if pid.nil?
-    # Shutting down server
+    return if pid.nil?
     Process.kill("KILL", pid) unless pid.nil?
     Process.wait(pid) unless pid.nil? # Ensure the process has exited and the port is released 
-    pid = nil       
   end
 
 end
diff --git a/spec/support/ui_support.rb b/spec/support/ui_support.rb
index dc0c7fe7e..129be11c6 100644
--- a/spec/support/ui_support.rb
+++ b/spec/support/ui_support.rb
@@ -10,29 +10,36 @@ require 'spec/support/constants.rb'
 def start_beef_and_hook_browser()
     reset_beef_db
     pid = start_beef_server_and_wait
-    beef_session = BeefTest.login
-    hooked_browser = BeefTest.new_victim
 
-    expect(hooked_browser).not_to be_nil
-    expect(hooked_browser).to be_a(Capybara::Session)
-    expect(hooked_browser).to have_content('BeEF', wait: PAGE_LOAD_TIMEOUT)
+    begin
+        beef_session = BeefTest.login
+        hooked_browser = BeefTest.new_victim
 
-    expect(beef_session).not_to be_nil
-    expect(beef_session).to be_a(Capybara::Session)
-    expect(beef_session).to have_content('Hooked Browsers', wait: PAGE_LOAD_TIMEOUT)
+        expect(hooked_browser).not_to be_nil
+        expect(hooked_browser).to be_a(Capybara::Session)
+        expect(hooked_browser).to have_content('BeEF', wait: PAGE_LOAD_TIMEOUT)
 
-    navigate_to_hooked_browser(beef_session)
+        expect(beef_session).not_to be_nil
+        expect(beef_session).to be_a(Capybara::Session)
+        expect(beef_session).to have_content('Hooked Browsers', wait: PAGE_LOAD_TIMEOUT)
 
-    expect(beef_session).to have_content('Commands', wait: PAGE_LOAD_TIMEOUT)
-    beef_session.click_on('Commands')
+        navigate_to_hooked_browser(beef_session)
 
-    return pid, beef_session, hooked_browser
+        expect(beef_session).to have_content('Commands', wait: PAGE_LOAD_TIMEOUT)
+        beef_session.click_on('Commands')
+
+        return pid, beef_session, hooked_browser
+    rescue => e
+        # If setup fails, cleanup the server before re-raising
+        stop_beef_server(pid)
+        raise e
+    end
 end
 
 def stop_beef_and_unhook_browser(pid, beef_session, hooked_browser)
     stop_beef_server(pid)
-    beef_session.driver.browser.close
-    hooked_browser.driver.browser.close
+    beef_session.driver.browser.close if beef_session
+    hooked_browser.driver.browser.close if hooked_browser
 end
 
 def navigate_to_hooked_browser(session, hooked_browser_text = nil)