from __future__ import annotations import ast import hashlib from io import StringIO import pickle import re import sys from collections.abc import Callable from datetime import date, datetime, timedelta from pathlib import Path from requests.cookies import RequestsCookieJar from app.core.db import get_db from app.core.redis_client import get_redis from app.core.config import settings from app.core.files import read_runtime_json, write_runtime_json from app.services.import_worker_service import normalize_domain DOMAIN_ROOT = Path(settings.domain_root) if str(DOMAIN_ROOT) not in sys.path: sys.path.insert(0, str(DOMAIN_ROOT)) from detect.juming import JM # type: ignore # noqa: E402 from detect.jucha import JC # type: ignore # noqa: E402 YKJ_DOMAIN_PATTERN = re.compile(r" dict: return { "mode": "delete_list", "page_start": 1, "page_size": 500, "page_count": 1, "crawl_date": date.today().isoformat(), "auto_date": True, } def get_juming_preferences() -> dict: defaults = _default_juming_preferences() stored = read_runtime_json(JUMING_PREFERENCES_FILE, default={}) payload = { "mode": str(stored.get("mode", defaults["mode"])) if stored else defaults["mode"], "page_start": int(stored.get("page_start", defaults["page_start"])) if stored else defaults["page_start"], "page_size": int(stored.get("page_size", defaults["page_size"])) if stored else defaults["page_size"], "page_count": int(stored.get("page_count", defaults["page_count"])) if stored else defaults["page_count"], "crawl_date": str(stored.get("crawl_date", defaults["crawl_date"])) if stored else defaults["crawl_date"], "auto_date": bool(stored.get("auto_date", defaults["auto_date"])) if stored else defaults["auto_date"], } if payload["mode"] not in {"delete_list", "fixed_price"}: payload["mode"] = defaults["mode"] payload["page_start"] = max(payload["page_start"], 1) payload["page_size"] = min(max(payload["page_size"], 1), 1000) payload["page_count"] = min(max(payload["page_count"], 1), 20) if not payload["crawl_date"]: payload["crawl_date"] = defaults["crawl_date"] return payload def update_juming_preferences(payload: dict) -> dict: current = get_juming_preferences() next_payload = { "mode": str(payload.get("mode", current["mode"]) or current["mode"]), "page_start": int(payload.get("page_start", current["page_start"]) or current["page_start"]), "page_size": int(payload.get("page_size", current["page_size"]) or current["page_size"]), "page_count": int(payload.get("page_count", current["page_count"]) or current["page_count"]), "crawl_date": str(payload.get("crawl_date", current["crawl_date"]) or current["crawl_date"]), "auto_date": bool(payload.get("auto_date", current["auto_date"])), } if next_payload["mode"] not in {"delete_list", "fixed_price"}: raise ValueError("无效的聚名采集类型") next_payload["page_start"] = max(next_payload["page_start"], 1) next_payload["page_size"] = min(max(next_payload["page_size"], 1), 1000) next_payload["page_count"] = min(max(next_payload["page_count"], 1), 20) write_runtime_json(JUMING_PREFERENCES_FILE, next_payload) return next_payload 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 _emit_log(log: Callable[[str], None] | None, message: str) -> None: if log: log(message) def _check_stop(should_stop: Callable[[], bool] | None) -> None: if should_stop and should_stop(): raise TaskStoppedError("任务已停止") 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_juming_cookie(cookie_jar: RequestsCookieJar) -> None: JUMING_COOKIE_FILE.parent.mkdir(parents=True, exist_ok=True) with JUMING_COOKIE_FILE.open("wb") as handle: pickle.dump(cookie_jar, handle) try: redis_client = get_redis() redis_client.set("domain_tool:juming_cookies", str(_cookie_jar_to_dict(cookie_jar))) except Exception: pass def _empty_import_stats() -> dict[str, int]: return { "total": 0, "valid": 0, "added": 0, "exists": 0, "invalid": 0, "failed": 0, } def _merge_import_stats(base: dict[str, int], delta: dict[str, int]) -> dict[str, int]: merged = dict(base or _empty_import_stats()) for key in ("total", "valid", "added", "exists", "invalid", "failed"): merged[key] = int(merged.get(key, 0) or 0) + int((delta or {}).get(key, 0) or 0) return merged def _load_delete_import_state() -> dict[str, dict]: payload = read_runtime_json(JUMING_DELETE_IMPORT_STATE_FILE, default={}) return payload if isinstance(payload, dict) else {} def _save_delete_import_state(payload: dict[str, dict]) -> None: write_runtime_json(JUMING_DELETE_IMPORT_STATE_FILE, payload) def _compute_domains_signature(domains: list[str]) -> str: digest = hashlib.sha1() for domain in domains: digest.update(str(domain).strip().encode("utf-8", errors="ignore")) digest.update(b"\n") return digest.hexdigest() def _looks_like_login_redirect(location: str) -> bool: normalized = str(location or "").strip().lower() if not normalized: return False return any( marker in normalized for marker in ( "/login", "user_zh", "p_login", "passport", "sign", ) ) def _looks_like_login_body(body: str) -> bool: normalized = str(body or "").strip().lower() if not normalized: return False return any( marker in normalized for marker in ( "账号登录", "请先登录", "登录后查看", "登录聚名", "user_zh", "p_login", ) ) def _validate_juming_cookie( cookie_jar: RequestsCookieJar | None, *, probe_date: str | None = None, ) -> tuple[bool, str]: if cookie_jar is None or not _cookie_jar_to_dict(cookie_jar): return False, "未检测到有效 Cookie" jm = JM() jm.cookie = cookie_jar probe_date = str(probe_date or date.today().isoformat()) url = f"{jm.base_url}/newcha/del_down?scsj={probe_date}" try: response = jm.session.get( url, headers=jm.headers, cookies=jm.cookie, allow_redirects=False, timeout=10, ) except Exception as exc: return False, f"登录态校验失败: {exc}" location = str(response.headers.get("Location") or "").strip() if response.status_code in {301, 302, 303, 307, 308}: if _looks_like_login_redirect(location): return False, "聚名登录态已失效,请重新登录" if location: return True, f"删除列表下载链路校验通过: {probe_date}" try: body = response.text except Exception: body = "" if _looks_like_login_body(body): return False, "聚名登录态已失效,请重新登录" if response.ok: return True, f"聚名 Cookie 已通过远端校验: {probe_date}" return False, f"聚名登录态校验失败,HTTP {response.status_code}" def _jucha_cookie_status() -> dict: if JUCHA_COOKIE_FILE.exists(): return { "cookie_ready": True, "cookie_file": str(JUCHA_COOKIE_FILE), } return { "cookie_ready": False, "cookie_file": str(JUCHA_COOKIE_FILE), } def _load_juming_cookie() -> tuple[RequestsCookieJar | None, str]: if JUMING_COOKIE_FILE.exists(): try: with JUMING_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 for legacy_path in LEGACY_JUMING_COOKIE_FILES: if not legacy_path.exists() or legacy_path == JUMING_COOKIE_FILE: continue try: with legacy_path.open("rb") as handle: loaded = pickle.load(handle) if isinstance(loaded, RequestsCookieJar): _persist_juming_cookie(loaded) return loaded, f"migrated:{legacy_path}" if isinstance(loaded, dict): cookie_jar = _cookie_dict_to_jar({str(k): str(v) for k, v in loaded.items()}) _persist_juming_cookie(cookie_jar) return cookie_jar, f"migrated:{legacy_path}" except Exception: continue try: redis_client = get_redis() raw = redis_client.get("domain_tool:juming_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 JUMING_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_juming_status() -> dict: cookie_jar, storage = _load_juming_cookie() cookie_valid, cookie_message = _validate_juming_cookie(cookie_jar) cookie_count = len(_cookie_jar_to_dict(cookie_jar)) if cookie_jar is not None else 0 status = { "cookie_ready": bool(cookie_jar is not None and cookie_valid), "cookie_present": cookie_jar is not None, "cookie_valid": cookie_valid, "cookie_message": cookie_message, "cookie_storage": storage, "cookie_file": str(JUMING_COOKIE_FILE), "cookie_count": cookie_count, "jucha": _jucha_cookie_status(), "supported_modes": [ {"label": "聚名一口价", "value": "fixed_price", "source_type": FIXED_PRICE_SOURCE_TYPE}, {"label": "聚名过期删除", "value": "delete_list", "source_type": DELETE_LIST_SOURCE_TYPE}, ], "defaults": get_juming_preferences(), } status["linked_jucha"] = { "attempted": False, "ok": bool(status["jucha"]["cookie_ready"]), "message": "聚查登录态已就绪" if status["jucha"]["cookie_ready"] else "尚未检测到聚查 Cookie", } return status def login_juming(email: str, password: str) -> dict: account = str(email or "").strip() secret = str(password or "").strip() if not account or not secret: raise ValueError("请输入聚名账号和密码") jm = JM() jm.load_cookies() login_result = jm.user_zh_p_login(account, secret) if not login_result[0]: raise ValueError(f"聚名登录失败: {login_result[1]}") jm.save_cookies() _persist_juming_cookie(jm.cookie) linked_jucha = { "attempted": True, "ok": False, "message": "未执行", } try: jc = JC() jc.load_juming_cookies() linked_ok, linked_message = jc.auth_login() linked_jucha["ok"] = bool(linked_ok) linked_jucha["message"] = str(linked_message) if linked_ok: jc.save_cookies() except Exception as exc: linked_jucha["message"] = f"聚查联名登录失败: {exc}" status = get_juming_status() status["linked_jucha"] = linked_jucha return status def login_jucha_with_juming_cookie() -> dict: if not JUMING_COOKIE_FILE.exists(): raise ValueError("请先完成聚名登录,当前未检测到聚名 Cookie") jc = JC() jc.load_juming_cookies() linked_ok, linked_message = jc.auth_login() if not linked_ok: raise ValueError(f"聚查登录失败: {linked_message}") jc.save_cookies() status = get_juming_status() status["linked_jucha"] = { "attempted": True, "ok": True, "message": str(linked_message), } return status def upload_juming_cookie(filename: str, content: bytes) -> dict: name = Path(filename or "juming_cookies.pkl").name.lower() if not (name.endswith(".pkl") or name.endswith(".pickle") or name.endswith(".json") or name.endswith(".txt")): raise ValueError("仅支持上传 .pkl / .pickle / .json / .txt 格式的聚名 Cookie 文件") cookie_jar: RequestsCookieJar | None = None parse_error: str | None = None if name.endswith(".pkl") or name.endswith(".pickle"): try: loaded = pickle.loads(content) if isinstance(loaded, RequestsCookieJar): cookie_jar = loaded elif isinstance(loaded, dict): cookie_jar = _cookie_dict_to_jar({str(k): str(v) for k, v in loaded.items()}) except Exception as exc: parse_error = str(exc) else: try: text = content.decode("utf-8") parsed = ast.literal_eval(text) if isinstance(parsed, dict): cookie_jar = _cookie_dict_to_jar({str(k): str(v) for k, v in parsed.items()}) except Exception as exc: parse_error = str(exc) if cookie_jar is None: raise ValueError(f"聚名 Cookie 文件解析失败: {parse_error or '内容不符合预期'}") if not _cookie_jar_to_dict(cookie_jar): raise ValueError("聚名 Cookie 文件为空,未检测到有效 Cookie") _persist_juming_cookie(cookie_jar) return { "cookie_ready": True, "cookie_storage": "upload", "cookie_file": str(JUMING_COOKIE_FILE), "cookie_count": len(_cookie_jar_to_dict(cookie_jar)), } def _insert_domains( domains: list[str], source_type: int, log: Callable[[str], None] | None = None, should_stop: Callable[[], bool] | None = None, *, announce_total: bool = True, progress_label: str = "", ) -> dict: total = len(domains) invalid = 0 progress_prefix = f"{progress_label} " if str(progress_label or "").strip() else "" if announce_total: _emit_log(log, f"{progress_prefix}开始入库处理,共收到 {total} 个原始域名") inserted = 0 existing = 0 processed = 0 valid = 0 last_progress_at = 0 pending_batch: list[tuple[str, str]] = [] pending_seen: set[str] = set() stage_ready = False def emit_progress(force: bool = False) -> None: nonlocal last_progress_at if not force and processed - last_progress_at < IMPORT_PROGRESS_EVERY: return last_progress_at = processed _emit_log( log, ( f"{progress_prefix}入库进度:已处理 {processed}/{total}," f"有效 {valid},新增 {inserted},已存在 {existing},无效 {invalid}" ), ) def ensure_stage_table(cur) -> None: nonlocal stage_ready if stage_ready: return cur.execute( """ create temporary table if not exists juming_import_stage ( domain text primary key, tld text not null ) on commit preserve rows """ ) stage_ready = True def stage_rows(cur, rows: list[tuple[str, str]]) -> None: buffer = StringIO() for domain, tld in rows: buffer.write(f"{domain}\t{tld}\n") buffer.seek(0) cur.copy_from(buffer, "juming_import_stage", columns=("domain", "tld")) def flush_batch(cur, conn) -> None: nonlocal inserted, existing if not pending_batch: return ensure_stage_table(cur) cur.execute("set local synchronous_commit = off") cur.execute("truncate table juming_import_stage") stage_rows(cur, pending_batch) cur.execute( """ with existing_rows as ( select count(*) from juming_import_stage stage join domains existing on existing.domain = stage.domain ), inserted as ( insert into domains ( domain, tld, source_type, use_status, detect_status, register_status, has_beian, company_type, website_url, beian_year, snapshot_years, expire_date, create_time, update_time, review_status, detect_time, backlink_count, jucha_status, juziseo_status ) select stage.domain, stage.tld, %s, 0, 0, 0, 1, null, null, null, null, null, now(), now(), 0, null, 0, 0, 0 from juming_import_stage stage left join domains existing on existing.domain = stage.domain where existing.id is null returning id ), task_insert as ( insert into detect_tasks ( domain_id, task_type, status, priority, retry_count, create_time, update_time ) select id, 1, 1, 5, 0, now(), now() from inserted returning 1 ) select (select count(*) from inserted), (select count(*) from task_insert), (select count(*) from existing_rows) """, (source_type,), ) inserted_count, _task_count, existing_count = cur.fetchone() inserted += int(inserted_count or 0) existing += int(existing_count or 0) conn.commit() pending_batch.clear() pending_seen.clear() with get_db() as conn: with conn.cursor() as cur: for value in domains: _check_stop(should_stop) processed += 1 normalized = normalize_domain(value) if not normalized: invalid += 1 emit_progress() continue valid += 1 if normalized in pending_seen: existing += 1 emit_progress() continue pending_seen.add(normalized) pending_batch.append((normalized, normalized.rsplit(".", 1)[-1])) if len(pending_batch) >= IMPORT_BATCH_SIZE: flush_batch(cur, conn) emit_progress(force=True) flush_batch(cur, conn) emit_progress(force=True) _emit_log(log, f"{progress_prefix}入库完成:有效 {valid},新增 {inserted},已存在 {existing},无效 {invalid}") return { "total": total, "valid": valid, "added": inserted, "exists": existing, "invalid": invalid, "failed": max(valid - existing - inserted, 0), } def _crawl_fixed_price( page_start: int, page_size: int, page_count: int, log: Callable[[str], None] | None = None, should_stop: Callable[[], bool] | None = None, ) -> tuple[list[str], list[dict[str, int]]]: cookie_jar, _ = _load_juming_cookie() jm = JM() jm.cookie = cookie_jar or RequestsCookieJar() domains: list[str] = [] pages: list[dict[str, int]] = [] current_page = page_start _emit_log(log, f"开始采集一口价域名:起始页 {page_start},每页 {page_size},最多 {page_count} 页") for _ in range(page_count): _check_stop(should_stop) _emit_log(log, f"正在抓取第 {current_page} 页") success, html = jm.ykj_get_list(page=current_page, page_size=page_size) if not success: raise RuntimeError(str(html)) page_domains = [item.strip() for item in YKJ_DOMAIN_PATTERN.findall(html) if item.strip()] domains.extend(page_domains) pages.append({"page": current_page, "count": len(page_domains)}) _emit_log(log, f"第 {current_page} 页抓取到 {len(page_domains)} 个域名,累计 {len(domains)} 个") if len(page_domains) < page_size: _emit_log(log, "当前页返回数量小于分页数量,判定已到末页,停止继续抓取") break current_page += 1 return domains, pages def _crawl_delete_list( crawl_date: str, auto_date: bool, log: Callable[[str], None] | None = None, should_stop: Callable[[], bool] | None = None, ) -> tuple[list[str], list[dict[str, int]]]: cookie_jar, _ = _load_juming_cookie() jm = JM() jm.cookie = cookie_jar or RequestsCookieJar() start_date = datetime.strptime(crawl_date, "%Y-%m-%d").date() end_date = date.today() + timedelta(days=4) current_date = start_date domains: list[str] = [] dates: list[dict[str, int]] = [] _emit_log(log, f"开始采集删除列表:起始日期 {crawl_date},自动追加日期 {'开启' if auto_date else '关闭'}") while current_date <= end_date: _check_stop(should_stop) _emit_log(log, f"正在抓取 {current_date.isoformat()} 的删除列表") domains_for_date = [item.strip() for item in jm.new_cha_del(current_date.isoformat()) if item.strip()] domains.extend(domains_for_date) dates.append({"date": current_date.isoformat(), "count": len(domains_for_date)}) _emit_log(log, f"{current_date.isoformat()} 抓取到 {len(domains_for_date)} 个域名,累计 {len(domains)} 个") if not auto_date: break current_date += timedelta(days=1) return domains, dates def _crawl_delete_list_and_import( crawl_date: str, auto_date: bool, log: Callable[[str], None] | None = None, should_stop: Callable[[], bool] | None = None, ) -> dict: cookie_jar, _ = _load_juming_cookie() jm = JM() jm.cookie = cookie_jar or RequestsCookieJar() start_date = datetime.strptime(crawl_date, "%Y-%m-%d").date() end_date = date.today() + timedelta(days=4) current_date = start_date dates: list[dict[str, int]] = [] domains_found = 0 sample_domains: list[str] = [] stats = _empty_import_stats() import_state = _load_delete_import_state() _emit_log(log, f"开始采集删除列表:起始日期 {crawl_date},自动追加日期 {'开启' if auto_date else '关闭'}") while current_date <= end_date: _check_stop(should_stop) current_date_text = current_date.isoformat() _emit_log(log, f"正在抓取 {current_date_text} 的删除列表") domains_for_date = [item.strip() for item in jm.new_cha_del(current_date_text) if item.strip()] domains_found += len(domains_for_date) dates.append({"date": current_date_text, "count": len(domains_for_date)}) _emit_log(log, f"{current_date_text} 抓取到 {len(domains_for_date)} 个域名,累计 {domains_found} 个") if domains_for_date: if len(sample_domains) < 20: sample_domains.extend(domains_for_date[: max(0, 20 - len(sample_domains))]) signature = _compute_domains_signature(domains_for_date) cached = import_state.get(current_date_text) or {} if ( cached.get("signature") == signature and int(cached.get("total", 0) or 0) == len(domains_for_date) ): cached_valid = int(cached.get("valid", 0) or 0) cached_invalid = int(cached.get("invalid", 0) or 0) date_stats = { "total": len(domains_for_date), "valid": cached_valid, "added": 0, "exists": cached_valid, "invalid": cached_invalid, "failed": 0, } _emit_log( log, ( f"{current_date_text} 删除列表内容未变化,跳过重复入库:" f"有效 {cached_valid},视为已存在 {cached_valid},无效 {cached_invalid}" ), ) else: date_stats = _insert_domains( domains_for_date, DELETE_LIST_SOURCE_TYPE, log=log, should_stop=should_stop, announce_total=False, progress_label=current_date_text, ) import_state[current_date_text] = { "signature": signature, "total": int(date_stats.get("total", 0) or 0), "valid": int(date_stats.get("valid", 0) or 0), "invalid": int(date_stats.get("invalid", 0) or 0), "updated_at": datetime.now().isoformat(sep=" ", timespec="seconds"), } _save_delete_import_state(import_state) stats = _merge_import_stats(stats, date_stats) if not auto_date: break current_date += timedelta(days=1) _emit_log( log, ( f"删除列表采集+入库完成:抓取 {domains_found} 个域名," f"新增 {stats['added']},已存在 {stats['exists']},无效 {stats['invalid']}" ), ) return { "mode": "delete_list", "dates": dates, "domains_found": domains_found, "stats": stats, "sample_domains": sample_domains[:20], } def crawl_juming( payload: dict, log: Callable[[str], None] | None = None, should_stop: Callable[[], bool] | None = None, ) -> dict: mode = str(payload.get("mode") or "delete_list").strip() if mode not in {"fixed_price", "delete_list"}: raise ValueError("仅支持 fixed_price 或 delete_list") cookie_jar, storage = _load_juming_cookie() if cookie_jar is None: raise ValueError("未找到聚名 Cookie,请先在桌面版系统设置完成聚名登录,或将 Cookie 同步到服务器") cookie_valid, cookie_message = _validate_juming_cookie(cookie_jar) if not cookie_valid: raise ValueError(cookie_message) _emit_log(log, f"检测到聚名登录态,来源:{storage},远端校验通过") _check_stop(should_stop) if mode == "fixed_price": page_start = max(int(payload.get("page_start") or 1), 1) page_size = min(max(int(payload.get("page_size") or 500), 1), 1000) page_count = min(max(int(payload.get("page_count") or 1), 1), 20) domains, pages = _crawl_fixed_price(page_start, page_size, page_count, log=log, should_stop=should_stop) stats = _insert_domains(domains, FIXED_PRICE_SOURCE_TYPE, log=log, should_stop=should_stop) return { "mode": mode, "cookie_storage": storage, "pages": pages, "domains_found": len(domains), "stats": stats, "sample_domains": domains[:20], } crawl_date = str(payload.get("crawl_date") or date.today().isoformat()) auto_date = bool(payload.get("auto_date", True)) result = _crawl_delete_list_and_import(crawl_date, auto_date, log=log, should_stop=should_stop) result["cookie_storage"] = storage return result