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