summaryrefslogtreecommitdiffstats
path: root/g4f/webdriver.py
diff options
context:
space:
mode:
Diffstat (limited to 'g4f/webdriver.py')
-rw-r--r--g4f/webdriver.py76
1 files changed, 53 insertions, 23 deletions
diff --git a/g4f/webdriver.py b/g4f/webdriver.py
index da283409..e5ecd8bf 100644
--- a/g4f/webdriver.py
+++ b/g4f/webdriver.py
@@ -1,5 +1,4 @@
from __future__ import annotations
-
from platformdirs import user_config_dir
from selenium.webdriver.remote.webdriver import WebDriver
from undetected_chromedriver import Chrome, ChromeOptions
@@ -21,7 +20,16 @@ def get_browser(
proxy: str = None,
options: ChromeOptions = None
) -> WebDriver:
- if user_data_dir == None:
+ """
+ Creates and returns a Chrome WebDriver with the specified options.
+
+ :param user_data_dir: Directory for user data. If None, uses default directory.
+ :param headless: Boolean indicating whether to run the browser in headless mode.
+ :param proxy: Proxy settings for the browser.
+ :param options: ChromeOptions object with specific browser options.
+ :return: An instance of WebDriver.
+ """
+ if user_data_dir is None:
user_data_dir = user_config_dir("g4f")
if user_data_dir and debug.logging:
print("Open browser with config dir:", user_data_dir)
@@ -39,36 +47,45 @@ def get_browser(
headless=headless
)
-def get_driver_cookies(driver: WebDriver):
- return dict([(cookie["name"], cookie["value"]) for cookie in driver.get_cookies()])
+def get_driver_cookies(driver: WebDriver) -> dict:
+ """
+ Retrieves cookies from the given WebDriver.
+
+ :param driver: WebDriver from which to retrieve cookies.
+ :return: A dictionary of cookies.
+ """
+ return {cookie["name"]: cookie["value"] for cookie in driver.get_cookies()}
def bypass_cloudflare(driver: WebDriver, url: str, timeout: int) -> None:
- # Open website
+ """
+ Attempts to bypass Cloudflare protection when accessing a URL using the provided WebDriver.
+
+ :param driver: The WebDriver to use.
+ :param url: URL to access.
+ :param timeout: Time in seconds to wait for the page to load.
+ """
driver.get(url)
- # Is cloudflare protection
if driver.find_element(By.TAG_NAME, "body").get_attribute("class") == "no-js":
if debug.logging:
print("Cloudflare protection detected:", url)
try:
- # Click button in iframe
- WebDriverWait(driver, 5).until(
- EC.presence_of_element_located((By.CSS_SELECTOR, "#turnstile-wrapper iframe"))
- )
driver.switch_to.frame(driver.find_element(By.CSS_SELECTOR, "#turnstile-wrapper iframe"))
WebDriverWait(driver, 5).until(
EC.presence_of_element_located((By.CSS_SELECTOR, "#challenge-stage input"))
- )
- driver.find_element(By.CSS_SELECTOR, "#challenge-stage input").click()
- except:
- pass
+ ).click()
+ except Exception as e:
+ if debug.logging:
+ print(f"Error bypassing Cloudflare: {e}")
finally:
driver.switch_to.default_content()
- # No cloudflare protection
WebDriverWait(driver, timeout).until(
EC.presence_of_element_located((By.CSS_SELECTOR, "body:not(.no-js)"))
)
-class WebDriverSession():
+class WebDriverSession:
+ """
+ Manages a Selenium WebDriver session, including handling of virtual displays and proxies.
+ """
def __init__(
self,
webdriver: WebDriver = None,
@@ -81,9 +98,7 @@ class WebDriverSession():
self.webdriver = webdriver
self.user_data_dir = user_data_dir
self.headless = headless
- self.virtual_display = None
- if has_pyvirtualdisplay and virtual_display:
- self.virtual_display = Display(size=(1920, 1080))
+ self.virtual_display = Display(size=(1920, 1080)) if has_pyvirtualdisplay and virtual_display else None
self.proxy = proxy
self.options = options
self.default_driver = None
@@ -94,8 +109,15 @@ class WebDriverSession():
headless: bool = False,
virtual_display: bool = False
) -> WebDriver:
- if user_data_dir == None:
- user_data_dir = self.user_data_dir
+ """
+ Reopens the WebDriver session with the specified parameters.
+
+ :param user_data_dir: Directory for user data.
+ :param headless: Boolean indicating whether to run the browser in headless mode.
+ :param virtual_display: Boolean indicating whether to use a virtual display.
+ :return: An instance of WebDriver.
+ """
+ user_data_dir = user_data_dir or self.user_data_dir
if self.default_driver:
self.default_driver.quit()
if not virtual_display and self.virtual_display:
@@ -105,6 +127,10 @@ class WebDriverSession():
return self.default_driver
def __enter__(self) -> WebDriver:
+ """
+ Context management method for entering a session.
+ :return: An instance of WebDriver.
+ """
if self.webdriver:
return self.webdriver
if self.virtual_display:
@@ -113,11 +139,15 @@ class WebDriverSession():
return self.default_driver
def __exit__(self, exc_type, exc_val, exc_tb):
+ """
+ Context management method for exiting a session. Closes and quits the WebDriver.
+ """
if self.default_driver:
try:
self.default_driver.close()
- except:
- pass
+ except Exception as e:
+ if debug.logging:
+ print(f"Error closing WebDriver: {e}")
self.default_driver.quit()
if self.virtual_display:
self.virtual_display.stop() \ No newline at end of file