# Licensed to the Software Freedom Conservancy (SFC) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The SFC licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from selenium.webdriver.common.bidi.common import command_builder
if TYPE_CHECKING:
from selenium.webdriver.remote.websocket_connection import WebSocketConnection
[docs]
class SameSite:
"""Represents the possible same site values for cookies."""
STRICT = "strict"
LAX = "lax"
NONE = "none"
DEFAULT = "default"
[docs]
class BytesValue:
"""Represents a bytes value."""
TYPE_BASE64 = "base64"
TYPE_STRING = "string"
def __init__(self, type: str, value: str):
self.type = type
self.value = value
[docs]
def to_dict(self) -> dict[str, str]:
"""Converts the BytesValue to a dictionary.
Returns:
A dictionary representation of the BytesValue.
"""
return {"type": self.type, "value": self.value}
[docs]
class Cookie:
"""Represents a cookie."""
def __init__(
self,
name: str,
value: BytesValue,
domain: str,
path: str | None = None,
size: int | None = None,
http_only: bool | None = None,
secure: bool | None = None,
same_site: str | None = None,
expiry: int | None = None,
):
self.name = name
self.value = value
self.domain = domain
self.path = path
self.size = size
self.http_only = http_only
self.secure = secure
self.same_site = same_site
self.expiry = expiry
[docs]
@classmethod
def from_dict(cls, data: dict[str, Any]) -> Cookie:
"""Creates a Cookie instance from a dictionary.
Args:
data: A dictionary containing the cookie information.
Returns:
A new instance of Cookie.
"""
# Validation for empty strings
name = data.get("name")
if not name:
raise ValueError("name is required and cannot be empty")
domain = data.get("domain")
if not domain:
raise ValueError("domain is required and cannot be empty")
value = BytesValue(data.get("value", {}).get("type"), data.get("value", {}).get("value"))
return cls(
name=str(name),
value=value,
domain=str(domain),
path=data.get("path"),
size=data.get("size"),
http_only=data.get("httpOnly"),
secure=data.get("secure"),
same_site=data.get("sameSite"),
expiry=data.get("expiry"),
)
[docs]
class CookieFilter:
"""Represents a filter for cookies."""
def __init__(
self,
name: str | None = None,
value: BytesValue | None = None,
domain: str | None = None,
path: str | None = None,
size: int | None = None,
http_only: bool | None = None,
secure: bool | None = None,
same_site: str | None = None,
expiry: int | None = None,
):
self.name = name
self.value = value
self.domain = domain
self.path = path
self.size = size
self.http_only = http_only
self.secure = secure
self.same_site = same_site
self.expiry = expiry
[docs]
def to_dict(self) -> dict[str, Any]:
"""Converts the CookieFilter to a dictionary.
Returns:
A dictionary representation of the CookieFilter.
"""
result: dict[str, Any] = {}
if self.name is not None:
result["name"] = self.name
if self.value is not None:
result["value"] = self.value.to_dict()
if self.domain is not None:
result["domain"] = self.domain
if self.path is not None:
result["path"] = self.path
if self.size is not None:
result["size"] = self.size
if self.http_only is not None:
result["httpOnly"] = self.http_only
if self.secure is not None:
result["secure"] = self.secure
if self.same_site is not None:
result["sameSite"] = self.same_site
if self.expiry is not None:
result["expiry"] = self.expiry
return result
[docs]
class PartitionKey:
"""Represents a storage partition key."""
def __init__(self, user_context: str | None = None, source_origin: str | None = None):
self.user_context = user_context
self.source_origin = source_origin
[docs]
@classmethod
def from_dict(cls, data: dict[str, Any]) -> PartitionKey:
"""Creates a PartitionKey instance from a dictionary.
Args:
data: A dictionary containing the partition key information.
Returns:
A new instance of PartitionKey.
"""
return cls(
user_context=data.get("userContext"),
source_origin=data.get("sourceOrigin"),
)
[docs]
class BrowsingContextPartitionDescriptor:
"""Represents a browsing context partition descriptor."""
def __init__(self, context: str):
self.type = "context"
self.context = context
[docs]
def to_dict(self) -> dict[str, str]:
"""Converts the BrowsingContextPartitionDescriptor to a dictionary.
Returns:
Dict: A dictionary representation of the BrowsingContextPartitionDescriptor.
"""
return {"type": self.type, "context": self.context}
[docs]
class StorageKeyPartitionDescriptor:
"""Represents a storage key partition descriptor."""
def __init__(self, user_context: str | None = None, source_origin: str | None = None):
self.type = "storageKey"
self.user_context = user_context
self.source_origin = source_origin
[docs]
def to_dict(self) -> dict[str, str]:
"""Converts the StorageKeyPartitionDescriptor to a dictionary.
Returns:
Dict: A dictionary representation of the StorageKeyPartitionDescriptor.
"""
result = {"type": self.type}
if self.user_context is not None:
result["userContext"] = self.user_context
if self.source_origin is not None:
result["sourceOrigin"] = self.source_origin
return result
[docs]
class PartialCookie:
"""Represents a partial cookie for setting."""
def __init__(
self,
name: str,
value: BytesValue,
domain: str,
path: str | None = None,
http_only: bool | None = None,
secure: bool | None = None,
same_site: str | None = None,
expiry: int | None = None,
):
self.name = name
self.value = value
self.domain = domain
self.path = path
self.http_only = http_only
self.secure = secure
self.same_site = same_site
self.expiry = expiry
[docs]
def to_dict(self) -> dict[str, Any]:
"""Converts the PartialCookie to a dictionary.
Returns:
-------
Dict: A dictionary representation of the PartialCookie.
"""
result: dict[str, Any] = {
"name": self.name,
"value": self.value.to_dict(),
"domain": self.domain,
}
if self.path is not None:
result["path"] = self.path
if self.http_only is not None:
result["httpOnly"] = self.http_only
if self.secure is not None:
result["secure"] = self.secure
if self.same_site is not None:
result["sameSite"] = self.same_site
if self.expiry is not None:
result["expiry"] = self.expiry
return result
[docs]
class GetCookiesResult:
"""Represents the result of a getCookies command."""
def __init__(self, cookies: list[Cookie], partition_key: PartitionKey):
self.cookies = cookies
self.partition_key = partition_key
[docs]
@classmethod
def from_dict(cls, data: dict[str, Any]) -> GetCookiesResult:
"""Creates a GetCookiesResult instance from a dictionary.
Args:
data: A dictionary containing the get cookies result information.
Returns:
A new instance of GetCookiesResult.
"""
cookies = [Cookie.from_dict(cookie) for cookie in data.get("cookies", [])]
partition_key = PartitionKey.from_dict(data.get("partitionKey", {}))
return cls(cookies=cookies, partition_key=partition_key)
[docs]
class SetCookieResult:
"""Represents the result of a setCookie command."""
def __init__(self, partition_key: PartitionKey):
self.partition_key = partition_key
[docs]
@classmethod
def from_dict(cls, data: dict[str, Any]) -> SetCookieResult:
"""Creates a SetCookieResult instance from a dictionary.
Args:
data: A dictionary containing the set cookie result information.
Returns:
A new instance of SetCookieResult.
"""
partition_key = PartitionKey.from_dict(data.get("partitionKey", {}))
return cls(partition_key=partition_key)
[docs]
class DeleteCookiesResult:
"""Represents the result of a deleteCookies command."""
def __init__(self, partition_key: PartitionKey):
self.partition_key = partition_key
[docs]
@classmethod
def from_dict(cls, data: dict[str, Any]) -> DeleteCookiesResult:
"""Creates a DeleteCookiesResult instance from a dictionary.
Args:
data: A dictionary containing the delete cookies result information.
Returns:
A new instance of DeleteCookiesResult.
"""
partition_key = PartitionKey.from_dict(data.get("partitionKey", {}))
return cls(partition_key=partition_key)
[docs]
class Storage:
"""BiDi implementation of the storage module."""
def __init__(self, conn: WebSocketConnection) -> None:
self.conn = conn
[docs]
def get_cookies(
self,
filter: CookieFilter | None = None,
partition: BrowsingContextPartitionDescriptor | StorageKeyPartitionDescriptor | None = None,
) -> GetCookiesResult:
"""Gets cookies matching the specified filter.
Args:
filter: Optional filter to specify which cookies to retrieve.
partition: Optional partition key to limit the scope of the operation.
Returns:
A GetCookiesResult containing the cookies and partition key.
Example:
result = await storage.get_cookies(
filter=CookieFilter(name="sessionId"),
partition=PartitionKey(...)
)
"""
params = {}
if filter is not None:
params["filter"] = filter.to_dict()
if partition is not None:
params["partition"] = partition.to_dict()
result = self.conn.execute(command_builder("storage.getCookies", params))
return GetCookiesResult.from_dict(result)
[docs]
def set_cookie(
self,
cookie: PartialCookie,
partition: BrowsingContextPartitionDescriptor | StorageKeyPartitionDescriptor | None = None,
) -> SetCookieResult:
"""Sets a cookie in the browser.
Args:
cookie: The cookie to set.
partition: Optional partition descriptor.
Returns:
The result of the set cookie command.
"""
params = {"cookie": cookie.to_dict()}
if partition is not None:
params["partition"] = partition.to_dict()
result = self.conn.execute(command_builder("storage.setCookie", params))
return SetCookieResult.from_dict(result)
[docs]
def delete_cookies(
self,
filter: CookieFilter | None = None,
partition: BrowsingContextPartitionDescriptor | StorageKeyPartitionDescriptor | None = None,
) -> DeleteCookiesResult:
"""Deletes cookies that match the given parameters.
Args:
filter: Optional filter to match cookies to delete.
partition: Optional partition descriptor.
Returns:
The result of the delete cookies command.
"""
params = {}
if filter is not None:
params["filter"] = filter.to_dict()
if partition is not None:
params["partition"] = partition.to_dict()
result = self.conn.execute(command_builder("storage.deleteCookies", params))
return DeleteCookiesResult.from_dict(result)