Update features/cookies.py

features/cookies.py

@@ -0,0 +1,45 @@
+"""
+Cookie management features
+"""
+import json
+import logging
+from browser.driver import get_driver, cleanup_driver
+
+logger = logging.getLogger(__name__)
+
+def get_cookies(url: str, use_persistent: bool = False) -> str:
+    """Get all cookies from current domain"""
+    driver = None
+    try:
+        driver = get_driver(url, use_persistent)
+        cookies = driver.get_cookies()
+        return json.dumps(cookies, indent=2)
+    except Exception as e:
+        logger.error(f"Error in get_cookies: {e}")
+        return f"Error: {e}"
+    finally:
+        cleanup_driver(driver, use_persistent)
+
+def set_cookies(url: str, cookies_json: str, use_persistent: bool = False) -> str:
+    """Set cookies from JSON string"""
+    driver = None
+    try:
+        driver = get_driver(url, use_persistent)
+        
+        # Parse and add cookies
+        cookies = json.loads(cookies_json)
+        for cookie in cookies:
+            # Remove expiry if present (can cause issues)
+            if 'expiry' in cookie:
+                del cookie['expiry']
+            driver.add_cookie(cookie)
+        
+        # Refresh page to apply cookies
+        driver.refresh()
+        
+        return f"Set {len(cookies)} cookies successfully"
+    except Exception as e:
+        logger.error(f"Error in set_cookies: {e}")
+        return f"Error: {e}"
+    finally:
+        cleanup_driver(driver, use_persistent)