refactor victim

web-security/level-10/victim

@@ -1,15 +1,16 @@
 #!/opt/pwn.college/python
 
-import contextlib
+import psutil
 import urllib
+import atexit
 import time
 import sys
 import os
 
 from selenium import webdriver
 from selenium.webdriver.firefox.options import Options as FirefoxOptions
 from selenium.webdriver.firefox.service import Service as FirefoxService
-from selenium.webdriver.common.by import By   
+from selenium.webdriver.common.by import By
 from selenium.webdriver.support.wait import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.common.exceptions import TimeoutException, WebDriverException
@@ -18,57 +19,30 @@ os.setuid(os.geteuid())
 os.environ.clear()
 os.environ["PATH"] = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
 
-@contextlib.contextmanager
-def run_browser():
-    options = FirefoxOptions()
-    options.add_argument("--headless")
-
-    # workaround for ubuntu
-    if os.path.exists("/snap/bin/geckodriver"):
-        service = FirefoxService(executable_path="/snap/bin/geckodriver", log_path="/dev/null")
-    else:
-        service = FirefoxService(log_path="/dev/null")
-    driver = webdriver.Firefox(service=service, options=options)
-
-    try:
-        yield driver
-    finally:
-        driver.quit()
-
-if len(sys.argv) <= 1:
-    print(f"Usage: {sys.argv[0]} URL")
-    sys.exit(1)
-
-url = sys.argv[1]
-url_parsed = urllib.parse.urlparse(url)
-try:
-    assert url_parsed.hostname == "challenge.localhost", "hostname should be 'challenge.localhost'"
-    assert url_parsed.port in {None, 80, 8080}, "port should be 80 or 8080"
-except AssertionError as e:
-    print(f"Invalid URL: {e}")
-    sys.exit(2)
-
-
-print("Visiting the URL!")
-with run_browser() as browser:
-    try:
-        browser.get(url)
-        print("URL loaded...")
-
-        print("Logging in as admin...")
-        login = { "username": "admin", "password": open("/flag").read().strip() if url_parsed.port != 8080 else "pwn.college{fake_flag}" }
-        for name, value in login.items():
-            field = browser.find_element(By.NAME, name)
-            field.send_keys(value)
-        submit_field = browser.find_element(By.NAME, "submit")
-        submit_field.submit()
-
-        WebDriverWait(browser, 10).until(EC.staleness_of(submit_field))
-        time.sleep(3)
-        print("Visited! Go check if the attack worked!")
-    except WebDriverException as e:
-        if "can%E2%80%99t%20establish%20a%20connection" in str(e):
-            print("Connection error! Is the service running?")
-        else:
-            print(f"Failure: {e}...")
-        sys.exit(4)
+options = FirefoxOptions()
+options.add_argument("--headless")
+service = FirefoxService(log_path="/dev/null")
+
+browser = webdriver.Firefox(service=service, options=options)
+atexit.register(browser.quit)
+
+open_ports = { s.laddr.port for s in psutil.net_connections(kind="inet") if s.status == 'LISTEN' }
+if 80 in open_ports:
+    print("Port 80 is open. Will connect to the service there.")
+    challenge_url = "http://challenge.localhost:80/"
+    password = open("/flag").read().strip()
+elif 8080 in open_ports:
+    print("Port 8080 is open. Will connect to the service there.")
+    challenge_url = "http://challenge.localhost:8080/"
+    password = "pwn.college{fake_flag}"
+
+print("Visiting http://challenge.localhost/")
+browser.get("http://challenge.localhost")
+
+print("Logging in as admin...")
+browser.find_element(By.NAME, "username").send_keys("admin")
+browser.find_element(By.NAME, "password").send_keys(password)
+browser.find_element(By.NAME, "submit").submit()
+
+time.sleep(5)
+print("Visited! Go check if the attack worked!")