diff options
Diffstat (limited to '')
| -rw-r--r-- | g4f/webdriver.py | 111 |
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 |