57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
from utils.logger import setup_logger
|
|
|
|
logger = setup_logger('CookieManager')
|
|
|
|
|
|
class CookieManager:
|
|
def __init__(self, cookie_file: str = '.cookies.json'):
|
|
self.cookie_file = Path(cookie_file)
|
|
self.cookies: Dict[str, str] = {}
|
|
|
|
def set_cookies(self, cookies: Dict[str, str]):
|
|
self.cookies = cookies
|
|
self._save_cookies()
|
|
|
|
def get_cookies(self) -> Dict[str, str]:
|
|
if not self.cookies:
|
|
self._load_cookies()
|
|
return self.cookies
|
|
|
|
def get_cookie_string(self) -> str:
|
|
cookies = self.get_cookies()
|
|
return '; '.join([f"{k}={v}" for k, v in cookies.items()])
|
|
|
|
def _save_cookies(self):
|
|
try:
|
|
with open(self.cookie_file, 'w', encoding='utf-8') as f:
|
|
json.dump(self.cookies, f, ensure_ascii=False, indent=2)
|
|
except Exception as e:
|
|
logger.error(f"Failed to save cookies: {e}")
|
|
|
|
def _load_cookies(self):
|
|
if not self.cookie_file.exists():
|
|
return
|
|
|
|
try:
|
|
with open(self.cookie_file, 'r', encoding='utf-8') as f:
|
|
self.cookies = json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load cookies: {e}")
|
|
|
|
def validate_cookies(self) -> bool:
|
|
required_keys = {'msToken', 'ttwid', 'odin_tt', 'passport_csrf_token'}
|
|
cookies = self.get_cookies()
|
|
missing = [key for key in required_keys if key not in cookies or not cookies.get(key)]
|
|
if missing:
|
|
logger.warning(f"Cookie validation failed, missing: {', '.join(missing)}")
|
|
return False
|
|
return True
|
|
|
|
def clear_cookies(self):
|
|
self.cookies = {}
|
|
if self.cookie_file.exists():
|
|
self.cookie_file.unlink()
|