summaryrefslogtreecommitdiffstats
path: root/g4f/webdriver.py
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--g4f/webdriver.py111
1 files changed, 88 insertions, 23 deletions
diff --git a/g4f/webdriver.py b/g4f/webdriver.py
index da283409..9a83215f 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,19 @@ def get_browser(
proxy: str = None,
options: ChromeOptions = None
) -> WebDriver:
- if user_data_dir == None:
+ """
+ Creates and returns a Chrome WebDriver with specified options.
+
+ Args:
+ user_data_dir (str, optional): Directory for user data. If None, uses default directory.
+ headless (bool, optional): Whether to run the browser in headless mode. Defaults to False.
+ proxy (str, optional): Proxy settings for the browser. Defaults to None.
+ options (ChromeOptions, optional): ChromeOptions object with specific browser options. Defaults to None.
+
+ Returns:
+ WebDriver: An instance of WebDriver configured with the specified options.
+ """
+ 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 +50,53 @@ 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 specified WebDriver.
+
+ Args:
+ driver (WebDriver): The WebDriver instance from which to retrieve cookies.
+
+ Returns:
+ dict: A dictionary containing cookies with their names as keys and values as cookie values.
+ """
+ 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.
+
+ Args:
+ driver (WebDriver): The WebDriver to use for accessing the URL.
+ url (str): The URL to access.
+ timeout (int): Time in seconds to wait for the page to load.
+
+ Raises:
+ Exception: If there is an error while bypassing Cloudflare or loading the page.
+ """
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,
@@ -78,12 +106,21 @@ class WebDriverSession():
proxy: str = None,
options: ChromeOptions = None
):
+ """
+ Initializes a new instance of the WebDriverSession.
+
+ Args:
+ webdriver (WebDriver, optional): A WebDriver instance for the session. Defaults to None.
+ user_data_dir (str, optional): Directory for user data. Defaults to None.
+ headless (bool, optional): Whether to run the browser in headless mode. Defaults to False.
+ virtual_display (bool, optional): Whether to use a virtual display. Defaults to False.
+ proxy (str, optional): Proxy settings for the browser. Defaults to None.
+ options (ChromeOptions, optional): ChromeOptions for the browser. Defaults to None.
+ """
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 +131,18 @@ 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 new settings.
+
+ Args:
+ user_data_dir (str, optional): Directory for user data. Defaults to current value.
+ headless (bool, optional): Whether to run the browser in headless mode. Defaults to current value.
+ virtual_display (bool, optional): Whether to use a virtual display. Defaults to current value.
+
+ Returns:
+ WebDriver: The reopened WebDriver instance.
+ """
+ user_data_dir = user_data_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 +152,12 @@ class WebDriverSession():
return self.default_driver
def __enter__(self) -> WebDriver:
+ """
+ Context management method for entering a session. Initializes and returns a WebDriver instance.
+
+ Returns:
+ WebDriver: An instance of WebDriver for this session.
+ """
if self.webdriver:
return self.webdriver
if self.virtual_display:
@@ -113,11 +166,23 @@ 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.
+
+ Args:
+ exc_type: Exception type.
+ exc_val: Exception value.
+ exc_tb: Exception traceback.
+
+ Note:
+ Closes the WebDriver and stops the virtual display if used.
+ """
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