cookie/cookie_cloud.py

"""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 applogger import get_logger

logger = get_logger(__name__)


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.api_url = api_url.rstrip('/')
        self.uuid = uuid
        self.password = password

    def _get_crypt_key(self) -> bytes:
        """生成加密密钥"""
        combined_string = f"{self.uuid}-{self.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_data(self) -> Tuple[Optional[Dict], str]:
        """下载并解密所有数据"""
        try:
            url = f"{self.api_url}/get/{self.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=30)
            
            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:
            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_data()
        
        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)
        
        return processed_cookies

    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)
        
        return matched_cookies


if __name__ == "__main__":
    cc = CookieCloud(
        api_url="https://movie-pilot.org/cookiecloud",
        uuid="hu6n2vcUmzpu7mqUN2rVCg",
        password="jtDM5dV9AyqXkZdQVeA9f6"
    )
    
    cookies = cc.get_cookies()
    logger.info(f"获取到 {len(cookies)} 个域名的 cookies")
    
    lmkbi_cookies = cc.get_cookies_for_domain('lmkbi.95155.com')
    logger.info(f"获取到 {len(lmkbi_cookies)} 个 lmkbi.95155.com 的 cookies")