256 lines
8.3 KiB
Python
256 lines
8.3 KiB
Python
"""Cookie Cloud 模块 - 从 Cookie Cloud 服务器获取并处理 cookie 数据"""
|
||
import json
|
||
import hashlib
|
||
import base64
|
||
from typing import Dict, List, Optional, Tuple
|
||
import urllib.request
|
||
import urllib.error
|
||
from Crypto.Cipher import AES
|
||
from dataclasses import dataclass
|
||
from logger import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class CookieCloudConfig:
|
||
"""CookieCloud 配置"""
|
||
api_url: str
|
||
uuid: str
|
||
password: str
|
||
timeout: int = 30
|
||
|
||
|
||
class CookieCloud:
|
||
"""Cookie Cloud 客户端"""
|
||
|
||
def __init__(self, api_url: str, uuid: str, password: str):
|
||
"""
|
||
初始化 Cookie Cloud 客户端
|
||
|
||
Args:
|
||
api_url: Cookie Cloud API 地址
|
||
uuid: 用户 UUID
|
||
password: 用户密码
|
||
"""
|
||
self.config = CookieCloudConfig(
|
||
api_url=api_url.rstrip('/'),
|
||
uuid=uuid,
|
||
password=password,
|
||
timeout=30
|
||
)
|
||
|
||
def _get_crypt_key(self) -> bytes:
|
||
"""生成加密密钥"""
|
||
combined_string = f"{self.config.uuid}-{self.config.password}"
|
||
return hashlib.md5(combined_string.encode('utf-8')).hexdigest()[:16].encode("utf-8")
|
||
|
||
def _bytes_to_key(self, data: bytes, salt: bytes, output: int = 48) -> bytes:
|
||
"""OpenSSL EVP_BytesToKey 密钥派生算法"""
|
||
assert len(salt) == 8, len(salt)
|
||
data += salt
|
||
key = hashlib.md5(data).digest()
|
||
final_key = key
|
||
while len(final_key) < output:
|
||
key = hashlib.md5(key + data).digest()
|
||
final_key += key
|
||
return final_key[:output]
|
||
|
||
def _decrypt(self, encrypted: str, passphrase: bytes) -> bytes:
|
||
"""解密数据"""
|
||
encrypted_bytes = base64.b64decode(encrypted)
|
||
assert encrypted_bytes.startswith(b"Salted__"), "Invalid encrypted data format"
|
||
salt = encrypted_bytes[8:16]
|
||
key_iv = self._bytes_to_key(passphrase, salt, 32 + 16)
|
||
key = key_iv[:32]
|
||
iv = key_iv[32:]
|
||
aes = AES.new(key, AES.MODE_CBC, iv)
|
||
decrypted_padded = aes.decrypt(encrypted_bytes[16:])
|
||
padding_length = decrypted_padded[-1]
|
||
if isinstance(padding_length, str):
|
||
padding_length = ord(padding_length)
|
||
return decrypted_padded[:-padding_length]
|
||
|
||
def _download_all(self) -> Tuple[Optional[Dict], str]:
|
||
"""下载所有cookie和local storage数据"""
|
||
try:
|
||
url = f"{self.config.api_url}/get/{self.config.uuid}"
|
||
request = urllib.request.Request(
|
||
url,
|
||
headers={
|
||
'Content-Type': 'application/json',
|
||
'User-Agent': 'CookieCloud-Client/1.0'
|
||
},
|
||
method='GET'
|
||
)
|
||
|
||
response = urllib.request.urlopen(request, timeout=self.config.timeout)
|
||
|
||
if response.status != 200:
|
||
return None, f"服务器返回错误状态码: {response.status}"
|
||
|
||
result = json.loads(response.read().decode('utf-8'))
|
||
|
||
if not result:
|
||
return None, "服务器返回数据为空"
|
||
|
||
encrypted = result.get("encrypted")
|
||
if not encrypted:
|
||
return None, "未获取到cookie密文"
|
||
|
||
crypt_key = self._get_crypt_key()
|
||
try:
|
||
decrypted_data = self._decrypt(encrypted, crypt_key)
|
||
result = json.loads(decrypted_data.decode("utf-8"))
|
||
except Exception as e:
|
||
return None, f"cookie解密失败: {str(e)}"
|
||
|
||
if not result:
|
||
return None, "cookie解密为空"
|
||
|
||
return result, ""
|
||
|
||
except urllib.error.HTTPError as e:
|
||
if e.code == 401:
|
||
return None, "认证失败，请检查用户名和密码"
|
||
else:
|
||
return None, f"HTTP错误: {e.code} {e.reason}"
|
||
except urllib.error.URLError as e:
|
||
return None, f"网络连接失败: {e.reason}"
|
||
except Exception as e:
|
||
return None, f"下载失败: {str(e)}"
|
||
|
||
def get_cookies(self) -> Dict[str, List[Dict]]:
|
||
"""
|
||
从 Cookie Cloud 服务器获取 cookie 数据
|
||
|
||
Returns:
|
||
按域名分组的 cookie 数据字典
|
||
格式: {"domain.com": [{"name": "cookie1", "value": "val1", ...}, ...], ...}
|
||
"""
|
||
data, error = self._download_all()
|
||
|
||
if error:
|
||
logger.error(error)
|
||
return {}
|
||
|
||
# 处理数据结构
|
||
cookie_data = {}
|
||
|
||
# 兼容直接按域名分组的格式
|
||
if isinstance(data, dict) and not data.get('cookie_data'):
|
||
cookie_data = data
|
||
# 兼容包含 cookie_data 的格式
|
||
elif isinstance(data, dict) and data.get('cookie_data'):
|
||
cookie_data = data.get('cookie_data', {})
|
||
|
||
# 处理 sameSite 字段
|
||
processed_cookies = {}
|
||
for domain, cookies in cookie_data.items():
|
||
if not cookies:
|
||
continue
|
||
processed_cookies[domain] = []
|
||
for cookie in cookies:
|
||
if cookie.get('sameSite') == 'unspecified':
|
||
cookie['sameSite'] = 'Lax'
|
||
processed_cookies[domain].append(cookie)
|
||
|
||
logger.info(f"获取到 {len(processed_cookies)} 个域名的 cookies")
|
||
return processed_cookies
|
||
|
||
def get_local_storage(self) -> Dict[str, Dict]:
|
||
"""
|
||
从 Cookie Cloud 服务器获取 local storage 数据
|
||
|
||
Returns:
|
||
按域名分组的 local storage 数据字典
|
||
格式: {"domain.com": {"key1": "value1", ...}, ...}
|
||
"""
|
||
data, error = self._download_all()
|
||
|
||
if error:
|
||
logger.error(error)
|
||
return {}
|
||
|
||
# 处理数据结构
|
||
local_storage_data = {}
|
||
|
||
if isinstance(data, dict) and data.get('local_storage_data'):
|
||
local_storage_data = data.get('local_storage_data', {})
|
||
|
||
logger.info(f"获取到 {len(local_storage_data)} 个域名的 local storage")
|
||
return local_storage_data
|
||
|
||
def get_cookies_for_domain(self, domain: str) -> List[Dict]:
|
||
"""
|
||
获取指定域名的 cookie
|
||
|
||
Args:
|
||
domain: 域名
|
||
|
||
Returns:
|
||
cookie 列表
|
||
"""
|
||
all_cookies = self.get_cookies()
|
||
|
||
matched_cookies = []
|
||
for cookie_domain, cookies in all_cookies.items():
|
||
clean_domain = cookie_domain.lstrip('.')
|
||
if domain in clean_domain or clean_domain in domain:
|
||
matched_cookies.extend(cookies)
|
||
|
||
logger.info(f"获取到 {len(matched_cookies)} 个 {domain} 的 cookies")
|
||
return matched_cookies
|
||
|
||
def get_local_storage_for_domain(self, domain: str) -> Dict:
|
||
"""
|
||
获取指定域名的 local storage
|
||
|
||
Args:
|
||
domain: 域名
|
||
|
||
Returns:
|
||
local storage 字典
|
||
"""
|
||
all_local_storage = self.get_local_storage()
|
||
|
||
for ls_domain, data in all_local_storage.items():
|
||
clean_domain = ls_domain.lstrip('.')
|
||
if domain in clean_domain or clean_domain in domain:
|
||
logger.info(f"获取到 {domain} 的 local storage")
|
||
return data
|
||
|
||
return {}
|
||
|
||
def get_all_data(self) -> Dict:
|
||
"""
|
||
获取所有数据（包括 cookies 和 local storage）
|
||
|
||
Returns:
|
||
包含所有数据的字典
|
||
"""
|
||
data, error = self._download_all()
|
||
|
||
if error:
|
||
logger.error(error)
|
||
return {}
|
||
|
||
result = {
|
||
'cookies': {},
|
||
'local_storage': {}
|
||
}
|
||
|
||
# 处理 cookies
|
||
if isinstance(data, dict):
|
||
if not data.get('cookie_data'):
|
||
result['cookies'] = data
|
||
elif data.get('cookie_data'):
|
||
result['cookies'] = data.get('cookie_data', {})
|
||
|
||
# 处理 local storage
|
||
if data.get('local_storage_data'):
|
||
result['local_storage'] = data.get('local_storage_data', {})
|
||
|
||
logger.info(f"获取到所有数据: {len(result['cookies'])} 个域名的 cookies, {len(result['local_storage'])} 个域名的 local storage")
|
||
return result
|