fix test cleanup and install firefox dependencies

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

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."
+      raise "Port #{@port} is already in use. Cannot start BeEF server."
-      exit
     end
     load_beef_extensions_and_modules
     
@@ -317,11 +316,9 @@ require 'socket'
   end
 
   def stop_beef_server(pid)
-    exit if pid.nil?
+    return if pid.nil?
-    # 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 
-    pid = nil       
   end
 
 end

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
+    begin
-    expect(hooked_browser).to be_a(Capybara::Session)
+        beef_session = BeefTest.login
-    expect(hooked_browser).to have_content('BeEF', wait: PAGE_LOAD_TIMEOUT)
+        hooked_browser = BeefTest.new_victim
 
-    expect(beef_session).not_to be_nil
+        expect(hooked_browser).not_to be_nil
-    expect(beef_session).to be_a(Capybara::Session)
+        expect(hooked_browser).to be_a(Capybara::Session)
-    expect(beef_session).to have_content('Hooked Browsers', wait: PAGE_LOAD_TIMEOUT)
+        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)
+        navigate_to_hooked_browser(beef_session)
-    beef_session.click_on('Commands')
 
-    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
+    beef_session.driver.browser.close if beef_session
-    hooked_browser.driver.browser.close
+    hooked_browser.driver.browser.close if hooked_browser
 end
 
 def navigate_to_hooked_browser(session, hooked_browser_text = nil)