76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
|
import http.cookies
|
||
|
from typing import Optional
|
||
|
|
||
|
"""
|
||
|
_cookiejar.py
|
||
|
websocket - WebSocket client library for Python
|
||
|
|
||
|
Copyright 2024 engn33r
|
||
|
|
||
|
Licensed 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.
|
||
|
"""
|
||
|
|
||
|
|
||
|
class SimpleCookieJar:
|
||
|
def __init__(self) -> None:
|
||
|
self.jar: dict = {}
|
||
|
|
||
|
def add(self, set_cookie: Optional[str]) -> None:
|
||
|
if set_cookie:
|
||
|
simple_cookie = http.cookies.SimpleCookie(set_cookie)
|
||
|
|
||
|
for v in simple_cookie.values():
|
||
|
if domain := v.get("domain"):
|
||
|
if not domain.startswith("."):
|
||
|
domain = f".{domain}"
|
||
|
cookie = (
|
||
|
self.jar.get(domain)
|
||
|
if self.jar.get(domain)
|
||
|
else http.cookies.SimpleCookie()
|
||
|
)
|
||
|
cookie.update(simple_cookie)
|
||
|
self.jar[domain.lower()] = cookie
|
||
|
|
||
|
def set(self, set_cookie: str) -> None:
|
||
|
if set_cookie:
|
||
|
simple_cookie = http.cookies.SimpleCookie(set_cookie)
|
||
|
|
||
|
for v in simple_cookie.values():
|
||
|
if domain := v.get("domain"):
|
||
|
if not domain.startswith("."):
|
||
|
domain = f".{domain}"
|
||
|
self.jar[domain.lower()] = simple_cookie
|
||
|
|
||
|
def get(self, host: str) -> str:
|
||
|
if not host:
|
||
|
return ""
|
||
|
|
||
|
cookies = []
|
||
|
for domain, _ in self.jar.items():
|
||
|
host = host.lower()
|
||
|
if host.endswith(domain) or host == domain[1:]:
|
||
|
cookies.append(self.jar.get(domain))
|
||
|
|
||
|
return "; ".join(
|
||
|
filter(
|
||
|
None,
|
||
|
sorted(
|
||
|
[
|
||
|
f"{k}={v.value}"
|
||
|
for cookie in filter(None, cookies)
|
||
|
for k, v in cookie.items()
|
||
|
]
|
||
|
),
|
||
|
)
|
||
|
)
|