mirror of
https://github.com/PaiGramTeam/SIMNet.git
synced 2024-11-23 06:49:37 +00:00
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
from http.cookiejar import CookieJar
|
|
from http.cookies import SimpleCookie
|
|
from typing import Optional
|
|
|
|
from httpx import Cookies as _Cookies
|
|
|
|
from simnet.utils.types import CookieTypes
|
|
|
|
__all__ = ("Cookies",)
|
|
|
|
|
|
class Cookies(_Cookies):
|
|
"""A wrapper around `httpx.Cookies` that provides additional functionality."""
|
|
|
|
jar: CookieJar
|
|
|
|
def __init__(self, cookies: Optional[CookieTypes] = None): # skipcq: PYL-W0231
|
|
self.jar = CookieJar()
|
|
if cookies is None or isinstance(cookies, dict):
|
|
if isinstance(cookies, dict):
|
|
for key, value in cookies.items():
|
|
if isinstance(value, str):
|
|
self.set(key, value)
|
|
else:
|
|
self.set(key, str(value))
|
|
elif isinstance(cookies, list):
|
|
for key, value in cookies:
|
|
self.set(key, value)
|
|
elif isinstance(cookies, Cookies):
|
|
for cookie in cookies.jar:
|
|
self.jar.set_cookie(cookie)
|
|
elif isinstance(cookies, str):
|
|
cookie = SimpleCookie(cookies)
|
|
for key, value in cookie.items():
|
|
self.set(key, value.value)
|
|
else:
|
|
self.jar = cookies # type: ignore
|
|
|
|
COOKIE_USER_ID_NAMES = ("ltuid", "account_id", "stuid", "ltuid_v2", "account_id_v2")
|
|
|
|
@property
|
|
def account_id(self) -> Optional[int]:
|
|
"""Return the user account ID if present in the cookies.
|
|
|
|
If one of the user ID cookies exists in the cookies, return its integer value.
|
|
Otherwise, return `None`.
|
|
|
|
Returns:
|
|
Optional[int]: The user account ID, or `None` if it is not present in the cookies.
|
|
"""
|
|
for name in self.COOKIE_USER_ID_NAMES:
|
|
value = self.get(name)
|
|
if value is not None:
|
|
return int(value)
|
|
return None
|
|
|
|
def get(
|
|
self,
|
|
name: str,
|
|
default: Optional[str] = None,
|
|
domain: Optional[str] = None,
|
|
path: Optional[str] = None,
|
|
) -> Optional[str]:
|
|
"""
|
|
Get a cookie by name. May optionally include domain and path
|
|
in order to specify exactly which cookie to retrieve.
|
|
"""
|
|
value = None
|
|
for cookie in self.jar:
|
|
if (
|
|
cookie.name == name
|
|
and domain is None
|
|
or cookie.domain == domain
|
|
and path is None
|
|
or cookie.path == path
|
|
and cookie.value
|
|
):
|
|
value = cookie.value
|
|
if value is None:
|
|
return default
|
|
return value
|