CoolFace
Apppublic

zhiyinya/AIStudioBuildWS

sourceHugging Faceupdated 10mo agoView on Hugging Face
0likes
cookie_handler.py147 linesDownload Raw Back to utils
1def convert_cookie_editor_to_playwright(cookies_from_editor, logger=None):2    """3    将从 Cookie-Editor 插件导出的 Cookie 列表转换为 Playwright 兼容的格式。4    """5    playwright_cookies = []6    allowed_keys = {'name', 'value', 'domain', 'path', 'expires', 'httpOnly', 'secure', 'sameSite'}7 8    for cookie in cookies_from_editor:9        pw_cookie = {}10        for key in ['name', 'value', 'domain', 'path', 'httpOnly', 'secure']:11            if key in cookie:12                pw_cookie[key] = cookie[key]13        if cookie.get('session', False):14            pw_cookie['expires'] = -115        elif 'expirationDate' in cookie:16            if cookie['expirationDate'] is not None:17                pw_cookie['expires'] = int(cookie['expirationDate'])18            else:19                pw_cookie['expires'] = -120 21        if 'sameSite' in cookie:22            same_site_value = str(cookie['sameSite']).lower()23            if same_site_value == 'no_restriction':24                pw_cookie['sameSite'] = 'None'25            elif same_site_value in ['lax', 'strict']:26                pw_cookie['sameSite'] = same_site_value.capitalize()27            elif same_site_value == 'unspecified':28                pw_cookie['sameSite'] = 'Lax'29 30        if all(key in pw_cookie for key in ['name', 'value', 'domain', 'path']):31            playwright_cookies.append(pw_cookie)32        else:33            if logger:34                logger.warning(f"跳过一个格式不完整的 Cookie: {cookie}")35 36    return playwright_cookies37 38 39def convert_kv_to_playwright(kv_string, default_domain=".google.com", logger=None):40    """41    将键值对格式的 Cookie 字符串转换为 Playwright 兼容的格式。42 43    Args:44        kv_string (str): 包含 Cookie 的键值对字符串,格式为 "name1=value1; name2=value2; ..."45        default_domain (str): 默认域名,默认为".google.com"46        logger: 日志记录器47 48    Returns:49        list: Playwright 兼容的 Cookie 列表50    """51    import re52 53    playwright_cookies = []54 55    # 按分号分割 Cookie56    cookie_pairs = kv_string.split(';')57 58    for pair in cookie_pairs:59        pair = pair.strip()  # 去除首尾空白字符60 61        if not pair:  # 跳过空字符串62            continue63 64        # 跳过无效的 Cookie(不包含等号)65        if '=' not in pair:66            if logger:67                logger.warning(f"跳过无效的 Cookie 格式: '{pair}'")68            continue69 70        # 分割name和value71        name, value = pair.split('=', 1)  # 只分割第一个等号72        name = name.strip()73        value = value.strip()74 75        if not name:  # 跳过空名称76            if logger:77                logger.warning(f"跳过空名称的 Cookie: '{pair}'")78            continue79 80        # 构造 Playwright 格式的 Cookie81        pw_cookie = {82            'name': name,83            'value': value,84            'domain': default_domain,85            'path': '/',86            'expires': -1,  # 默认为会话 Cookie87            'httpOnly': False,  # KV 格式无法确定 httpOnly 状态,默认为 False88            'secure': True,     # 假设为安全 Cookie89            'sameSite': 'Lax'   # 默认 SameSite 策略90        }91 92        playwright_cookies.append(pw_cookie)93 94        if logger:95            logger.debug(f"成功转换 Cookie: {name} -> domain={default_domain}")96 97    return playwright_cookies98 99 100def auto_convert_to_playwright(cookie_data, default_domain=".google.com", logger=None):101    """102    自动识别 Cookie 数据格式并转换为 Playwright 兼容格式。103    支持两种输入格式:104    1. JSON 数组 (Cookie-Editor 导出格式)105    2. KV 字符串 (键值对格式: "name1=value1; name2=value2; ...")106 107    Args:108        cookie_data: Cookie 数据,可以是 list (JSON格式) 或 str (KV格式)109        default_domain (str): KV格式使用的默认域名,默认为".google.com"110        logger: 日志记录器111 112    Returns:113        list: Playwright 兼容的 Cookie 列表114 115    Raises:116        ValueError: 当格式无法识别时抛出异常117    """118    # 格式1: JSON 数组格式 (Cookie-Editor 导出格式)119    if isinstance(cookie_data, list):120        if logger:121            logger.debug(f"检测到 JSON 数组格式的 Cookie 数据,共 {len(cookie_data)} 个条目")122        return convert_cookie_editor_to_playwright(cookie_data, logger=logger)123 124    # 格式2: KV 字符串格式125    if isinstance(cookie_data, str):126        # 去除首尾空白字符127        cookie_str = cookie_data.strip()128 129        if not cookie_str:130            if logger:131                logger.warning("收到空的 Cookie 字符串")132            return []133 134        if logger:135            logger.debug(f"检测到 KV 字符串格式的 Cookie 数据")136 137        return convert_kv_to_playwright(138            cookie_str,139            default_domain=default_domain,140            logger=logger141        )142 143    # 无法识别的格式144    error_msg = f"无法识别的 Cookie 数据格式: {type(cookie_data).__name__}"145    if logger:146        logger.error(error_msg)147    raise ValueError(error_msg)