112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import ast
|
|
import pickle
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from requests.cookies import RequestsCookieJar
|
|
|
|
from app.core.config import settings
|
|
from app.core.redis_client import get_redis
|
|
|
|
|
|
DOMAIN_ROOT = Path(settings.domain_root)
|
|
if str(DOMAIN_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(DOMAIN_ROOT))
|
|
|
|
from detect.juziseo import Juziseo # type: ignore # noqa: E402
|
|
|
|
|
|
JUZISEO_COOKIE_FILE = DOMAIN_ROOT / "juziseo_cookies.pkl"
|
|
|
|
|
|
def _cookie_dict_to_jar(cookie_dict: dict[str, str]) -> RequestsCookieJar:
|
|
cookie_jar = RequestsCookieJar()
|
|
for name, value in cookie_dict.items():
|
|
cookie_jar.set(name, value)
|
|
return cookie_jar
|
|
|
|
|
|
def _cookie_jar_to_dict(cookie_jar: RequestsCookieJar) -> dict[str, str]:
|
|
return {str(cookie.name): str(cookie.value) for cookie in cookie_jar}
|
|
|
|
|
|
def _persist_juziseo_cookie(cookie_jar: RequestsCookieJar) -> None:
|
|
JUZISEO_COOKIE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
with JUZISEO_COOKIE_FILE.open("wb") as handle:
|
|
pickle.dump(cookie_jar, handle)
|
|
|
|
try:
|
|
redis_client = get_redis()
|
|
redis_client.set("domain_tool:juziseo_cookies", str(_cookie_jar_to_dict(cookie_jar)))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _load_juziseo_cookie() -> tuple[RequestsCookieJar | None, str]:
|
|
if JUZISEO_COOKIE_FILE.exists():
|
|
try:
|
|
with JUZISEO_COOKIE_FILE.open("rb") as handle:
|
|
loaded = pickle.load(handle)
|
|
if isinstance(loaded, RequestsCookieJar):
|
|
return loaded, "local"
|
|
if isinstance(loaded, dict):
|
|
return _cookie_dict_to_jar({str(k): str(v) for k, v in loaded.items()}), "local"
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
redis_client = get_redis()
|
|
raw = redis_client.get("domain_tool:juziseo_cookies")
|
|
if raw:
|
|
parsed = ast.literal_eval(raw)
|
|
if isinstance(parsed, dict) and parsed:
|
|
cookie_jar = _cookie_dict_to_jar({str(k): str(v) for k, v in parsed.items()})
|
|
try:
|
|
with JUZISEO_COOKIE_FILE.open("wb") as handle:
|
|
pickle.dump(cookie_jar, handle)
|
|
except Exception:
|
|
pass
|
|
return cookie_jar, "redis"
|
|
except Exception:
|
|
pass
|
|
|
|
return None, "missing"
|
|
|
|
|
|
def get_juziseo_status() -> dict:
|
|
cookie_jar, storage = _load_juziseo_cookie()
|
|
return {
|
|
"cookie_ready": cookie_jar is not None,
|
|
"cookie_storage": storage,
|
|
"cookie_file": str(JUZISEO_COOKIE_FILE),
|
|
"cookie_count": len(_cookie_jar_to_dict(cookie_jar)) if cookie_jar is not None else 0,
|
|
}
|
|
|
|
|
|
def login_juziseo(email: str, password: str) -> dict:
|
|
account = str(email or "").strip()
|
|
secret = str(password or "").strip()
|
|
if not account or not secret:
|
|
raise ValueError("请输入桔子SEO账号和密码")
|
|
|
|
juziseo = Juziseo()
|
|
juziseo.load_cookies(str(JUZISEO_COOKIE_FILE))
|
|
ok, message = juziseo.login(account, secret)
|
|
if not ok:
|
|
raise ValueError(f"桔子SEO登录失败: {message}")
|
|
|
|
try:
|
|
juziseo.save_cookies(str(JUZISEO_COOKIE_FILE))
|
|
except Exception:
|
|
pass
|
|
try:
|
|
_persist_juziseo_cookie(juziseo.cookie)
|
|
except Exception:
|
|
pass
|
|
|
|
status = get_juziseo_status()
|
|
status["message"] = str(message)
|
|
return status
|