yunke/yunke-server/scripts/scrape_via_page.py

#!/usr/bin/env python
"""页面浏览器抓取器 — 当 API 接口被风控时，用 Playwright 浏览器方式抓取数据。

原理：SPA 页面加载时会自己调用 friend_list/es 等接口渲染表格。
我们拦截页面的网络响应，获取原始 JSON —— 效果等同 API 调用，
但因为是浏览器页面上下文（非脚本直调），不触发风控。

用法：
    cd yunke-server
    DB_URL="sqlite+aiosqlite:///./smoke.db" .venv/Scripts/python.exe scripts/scrape_via_page.py --entity friends
    DB_URL="sqlite+aiosqlite:///./smoke.db" .venv/Scripts/python.exe scripts/scrape_via_page.py --entity groups
    DB_URL="sqlite+aiosqlite:///./smoke.db" .venv/Scripts/python.exe scripts/scrape_via_page.py --entity accounts

Cookie 从 .env 的 YK_COOKIE 读取。
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
from urllib.parse import urlparse

logger = logging.getLogger("scrape_via_page")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

YUNKE_BASE = "https://crm2.yunkecn.com"

# API endpoints we intercept (the SPA calls these to render tables)
INTERCEPT_PATHS = {
    "friend_list": "/wechatweb/wxFriends/friend_list/es",  # friends + groups (friendTypeYunke)
    "wxUser_list": "/wechatweb/wxUser/list",               # wechat accounts
    "chatroom": "/wechatweb/chatroom/getChatroomList",      # legacy groups
    "voice_list": "/wechatweb/call/list",                    # voice clips
    "call_list": "/wechatweb/callRecord/queryWeChatVoiceList",  # voice calls
    "wa_friend": "/qqweb/WhatsAppUser/businessFriendList",  # WA friends + groups
    "wa_user": "/qqweb/WhatsAppUser/businessSalesList",     # WA accounts
    "msg_chats": "/wechatweb/msg/chats",                    # chat messages
}


def parse_cookies(cookie_str: str, domain: str = "crm2.yunkecn.com") -> list[dict]:
    """Parse a cookie string like 'user=xxx; userToken=yyy' into Playwright cookie dicts."""
    cookies = []
    for pair in cookie_str.split(";"):
        pair = pair.strip()
        if "=" not in pair:
            continue
        name, value = pair.split("=", 1)
        cookies.append({
            "name": name.strip(),
            "value": value.strip(),
            "domain": domain,
            "path": "/",
        })
    return cookies


class PageScraper:
    """Navigates 云客 pages in a browser, intercepts API responses, extracts JSON."""

    def __init__(self, cookie_str: str, headless: bool = True):
        self.cookie_str = cookie_str
        self.headless = headless
        self.browser = None
        self.context = None
        self.page = None
        self._captured: list[dict] = []
        self._capture_path: str | None = None

    async def __aenter__(self):
        from playwright.async_api import async_playwright
        self._pw = await async_playwright().start()
        self.browser = await self._pw.chromium.launch(headless=self.headless)
        self.context = await self.browser.new_context(
            user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36",
            viewport={"width": 1920, "height": 1080},
        )
        # Set cookies
        cookies = parse_cookies(self.cookie_str)
        await self.context.add_cookies(cookies)
        self.page = await self.context.new_page()

        # Intercept API responses
        self.page.on("response", self._on_response)
        return self

    async def __aexit__(self, *args):
        if self.browser:
            await self.browser.close()
        if self._pw:
            await self._pw.stop()

    async def _on_response(self, response):
        """Capture API JSON responses that match our intercept paths."""
        url = response.url
        path = urlparse(url).path
        if self._capture_path and path == self._capture_path:
            try:
                body = await response.json()
                self._captured.append({"url": url, "status": response.status, "body": body})
                logger.info(f"  captured {path}: status={response.status}, keys={list(body.keys()) if isinstance(body, dict) else 'N/A'}")
            except Exception as e:
                logger.warning(f"  capture failed for {path}: {e}")

    def start_capture(self, api_path: str):
        """Start capturing responses for a specific API path."""
        self._capture_path = api_path
        self._captured.clear()

    def get_captured(self) -> list[dict]:
        """Return captured responses and stop capturing."""
        self._capture_path = None
        return self._captured

    async def navigate_and_wait(self, url: str, wait_selector: str = ".ant-table-row", timeout: int = 15000):
        """Navigate to a URL and wait for table rows to appear."""
        await self.page.goto(url, wait_until="domcontentloaded", timeout=timeout)
        try:
            await self.page.wait_for_selector(wait_selector, timeout=timeout)
        except Exception:
            pass  # might be empty page
        await self.page.wait_for_timeout(2000)  # extra settle time

    async def get_total_pages(self) -> int:
        """Extract total page count from Ant Design pagination."""
        try:
            total_text = await self.page.evaluate("""() => {
                const el = document.querySelector('.ant-pagination-total-text');
                return el ? el.textContent : '';
            }""")
            if total_text:
                import re
                m = re.search(r'(\d+)', total_text)
                if m:
                    total = int(m.group(1))
                    page_size = 10  # 云客 default
                    return (total + page_size - 1) // page_size
        except Exception:
            pass
        return 1

    async def click_next_page(self) -> bool:
        """Click the 'next page' button. Returns True if successful."""
        try:
            next_btn = self.page.locator('.ant-pagination-next:not(.ant-pagination-disabled)')
            if await next_btn.count() > 0:
                await next_btn.click()
                await self.page.wait_for_timeout(2000)  # wait for data load
                return True
        except Exception:
            pass
        return False

    async def scrape_entity_pages(
        self,
        list_url: str,
        api_path: str,
        max_pages: int = 100,
    ) -> list[dict]:
        """Navigate to a list page, paginate through all pages, capturing API responses.

        Returns a list of captured JSON response bodies (one per page).
        """
        self.start_capture(api_path)
        await self.navigate_and_wait(list_url)
        await self.page.wait_for_timeout(1000)  # let API call complete

        all_responses = list(self.get_captured())

        total_pages = await self.get_total_pages()
        logger.info(f"  total pages: {min(total_pages, max_pages)}")

        page_num = 1
        while page_num < min(total_pages, max_pages):
            self.start_capture(api_path)
            clicked = await self.click_next_page()
            if not clicked:
                break
            await self.page.wait_for_timeout(1500)
            all_responses.extend(self.get_captured())
            page_num += 1
            if page_num % 10 == 0:
                logger.info(f"  page {page_num}/{min(total_pages, max_pages)}")

        logger.info(f"  captured {len(all_responses)} API responses across {page_num} pages")
        return all_responses

    async def get_employee_links(self) -> list[dict]:
        """Extract employee wxIds + friend-count links from the staff list page."""
        await self.navigate_and_wait(f"{YUNKE_BASE}/cms/wechat/wechatlist")

        employees = await self.page.evaluate("""() => {
            const rows = document.querySelectorAll('.ant-table-row');
            const result = [];
            for (const tr of rows) {
                // Find the friend-count link (href contains wechatdetail/1/)
                const friendLink = tr.querySelector('a[href*="wechatdetail/1/"]');
                const groupLink = tr.querySelector('a[href*="wechatdetail/2/"]');
                if (friendLink) {
                    const href = friendLink.getAttribute('href');
                    const wxId = href.match(/wechatdetail\\/1\\/([^?]+)/)?.[1];
                    const friendCount = parseInt(friendLink.textContent) || 0;
                    const groupHref = groupLink?.getAttribute('href');
                    const groupCount = parseInt(groupLink?.textContent) || 0;
                    result.push({ wxId, friendCount, groupCount, friendHref: href, groupHref });
                }
            }
            return result;
        }""")

        # Paginate through ALL staff list pages
        all_employees = list(employees)
        page = 1
        while await self.click_next_page():
            page += 1
            await self.page.wait_for_timeout(2000)
            more = await self.page.evaluate("""() => {
                const rows = document.querySelectorAll('.ant-table-row');
                const result = [];
                for (const tr of rows) {
                    const friendLink = tr.querySelector('a[href*="wechatdetail/1/"]');
                    const groupLink = tr.querySelector('a[href*="wechatdetail/2/"]');
                    if (friendLink) {
                        const href = friendLink.getAttribute('href');
                        const wxId = href.match(/wechatdetail\\/1\\/([^?]+)/)?.[1];
                        const friendCount = parseInt(friendLink.textContent) || 0;
                        const groupHref = groupLink?.getAttribute('href');
                        const groupCount = parseInt(groupLink?.textContent) || 0;
                        result.push({ wxId, friendCount, groupCount, friendHref: href, groupHref });
                    }
                }
                return result;
            }""")
            all_employees.extend(more)
            logger.info(f"  staff page {page}: +{len(more)} employees (total: {len(all_employees)})")

        # Dedup by wxId
        seen = set()
        deduped = []
        for emp in all_employees:
            if emp["wxId"] not in seen:
                seen.add(emp["wxId"])
                deduped.append(emp)

        logger.info(f"Found {len(deduped)} unique employees with friend links ({page} pages)")
        return deduped


async def scrape_friends(cookie_str: str, max_employees: int = 100):
    """Scrape wechat customers (friends) via page navigation."""
    from app.db.base import async_session_maker
    from app.storage.repos import upsert_customer
    from app.scraper.wechat import _normalize  # reuse the existing normalizer

    total_fetched = 0
    total_upserted = 0

    async with PageScraper(cookie_str) as scraper:
        employees = await scraper.get_employee_links()

        for i, emp in enumerate(employees[:max_employees]):
            if emp["friendCount"] == 0:
                continue
            wx_id = emp["wxId"]
            url = f"{YUNKE_BASE}/cms/wechat/wechatlist/wechatdetail/1/{wx_id}?ymdStart=0000-00-00&ymdEnd=2099-12-31"
            logger.info(f"[{i+1}/{len(employees)}] {wx_id}: {emp['friendCount']} friends")

            responses = await scraper.scrape_entity_pages(
                list_url=url,
                api_path=INTERCEPT_PATHS["friend_list"],
                max_pages=200,
            )

            # Extract friend records from captured responses
            records = []
            for resp in responses:
                body = resp.get("body", {})
                data = body.get("data", {})
                rows = data.get("basicResult", {}).get("data", []) or data.get("data", [])
                for row in rows:
                    if isinstance(row, dict):
                        normalized = _normalize(row)
                        normalized["channel"] = "wechat"
                        normalized["owner_account_id"] = wx_id
                        records.append(normalized)

            total_fetched += len(records)

            # Batch upsert
            if records:
                async with async_session_maker() as session:
                    for rec in records:
                        await upsert_customer(session, rec)
                    await session.commit()
                total_upserted += len(records)
                logger.info(f"  upserted {len(records)} customers (total: {total_upserted})")

            # Human-like delay between employees
            await asyncio.sleep(3 + (i % 3))  # 3-5s random-ish

    logger.info(f"DONE: friends fetched={total_fetched}, upserted={total_upserted}")
    return total_fetched, total_upserted


async def scrape_groups(cookie_str: str, max_employees: int = 100):
    """Scrape wechat groups via page navigation."""
    from app.db.base import async_session_maker
    from app.storage.repos import upsert_chat_group
    from app.scraper.wechat import _normalize_chat_group

    total_fetched = 0
    total_upserted = 0

    async with PageScraper(cookie_str) as scraper:
        employees = await scraper.get_employee_links()

        for i, emp in enumerate(employees[:max_employees]):
            if emp.get("groupCount", 0) == 0 or not emp.get("groupHref"):
                continue
            wx_id = emp["wxId"]
            url = f"{YUNKE_BASE}{emp['groupHref']}"
            logger.info(f"[{i+1}/{len(employees)}] {wx_id}: {emp['groupCount']} groups")

            responses = await scraper.scrape_entity_pages(
                list_url=url,
                api_path=INTERCEPT_PATHS["friend_list"],
                max_pages=100,
            )

            records = []
            for resp in responses:
                body = resp.get("body", {})
                data = body.get("data", {})
                rows = data.get("basicResult", {}).get("data", []) or data.get("data", [])
                for row in rows:
                    if isinstance(row, dict):
                        normalized = _normalize_chat_group(row)
                        normalized["channel"] = "wechat"
                        records.append(normalized)

            total_fetched += len(records)
            if records:
                async with async_session_maker() as session:
                    for rec in records:
                        await upsert_chat_group(session, rec)
                    await session.commit()
                total_upserted += len(records)
                logger.info(f"  upserted {len(records)} groups (total: {total_upserted})")

            await asyncio.sleep(3 + (i % 3))

    logger.info(f"DONE: groups fetched={total_fetched}, upserted={total_upserted}")
    return total_fetched, total_upserted


async def scrape_messages(cookie_str: str, max_conversations: int = 100):
    """Scrape wechat messages via page navigation (chat window).

    Navigates to each conversation's chat window. The SPA calls /wechatweb/msg/chats
    to load messages — we intercept that response.
    """
    from app.db.base import async_session_maker, async_engine
    from sqlalchemy import text as sa_text
    from app.storage.repos import upsert_messages

    # Discover conversations from local customer table
    async with async_engine.connect() as conn:
        result = await conn.execute(sa_text(
            "SELECT DISTINCT owner_account_id, friend_id FROM customer "
            "WHERE channel='wechat' AND owner_account_id IS NOT NULL AND friend_id IS NOT NULL "
            "ORDER BY owner_account_id, friend_id"
        ))
        conversations = [(row[0], row[1]) for row in result.fetchall()]

    logger.info(f"Discovered {len(conversations)} conversations, will scrape {min(max_conversations, len(conversations))}")

    total_fetched = 0
    total_upserted = 0

    async with PageScraper(cookie_str) as scraper:
        page = scraper.page
        # First navigate to the CRM to establish session context
        await page.goto(f"{YUNKE_BASE}/cms/wechat/wechatlist", wait_until="domcontentloaded")
        await page.wait_for_timeout(3000)

        for i, (owner_id, friend_id) in enumerate(conversations[:max_conversations]):
            try:
                # Use browser fetch() to call the message API from the page context
                # (the browser sends it with proper cookies/headers → no 风控)
                body = await page.evaluate("""async ([ownerId, friendId]) => {
                    const resp = await fetch('/wechatweb/msg/chats', {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({
                            pageNum: 1,
                            pageSize: 50,
                            wechatId: ownerId,
                            friendId: friendId,
                            type: '',
                        }),
                    });
                    return await resp.json();
                }""", [owner_id, friend_id])

                responses = [{"body": body, "status": 200}]

            except Exception as e:
                logger.warning(f"  [{i+1}] {owner_id[:12]}:{friend_id[:12]} failed: {e}")
                # Re-navigate to reset the page if session expired
                await page.goto(f"{YUNKE_BASE}/cms/wechat/wechatlist", wait_until="domcontentloaded")
                await page.wait_for_timeout(2000)
                continue

            # Extract message records
            records = []
            for resp in responses:
                body = resp.get("body", {})
                data = body.get("data") or {}
                rows = data.get("data") or []
                for row in rows:
                    if not isinstance(row, dict):
                        continue
                    # Normalize to Message shape
                    ts_raw = row.get("createTime") or row.get("timestamp")
                    try:
                        ts = int(str(ts_raw)) if ts_raw else None
                    except (ValueError, TypeError):
                        ts = None
                    records.append({
                        "msg_id": str(row.get("msgId") or row.get("msgSvrId") or row.get("id") or f"{owner_id}:{friend_id}:{i}:{len(records)}"),
                        "channel": "wechat",
                        "owner_account_id": owner_id,
                        "talker_id": row.get("talker") or row.get("oriTalker") or friend_id,
                        "room_id": None,
                        "is_send": bool(row.get("isSend", 0)) if row.get("isSend") is not None else None,
                        "type": _normalize_msg_type(row.get("type")),
                        "type_raw": row.get("type"),
                        "text": row.get("content") or row.get("text") or "",
                        "file_url": row.get("file") or row.get("fileUrl"),
                        "timestamp": ts,
                        "sender_name": row.get("name") or row.get("senderName"),
                        "raw_json": row,
                    })

            total_fetched += len(records)

            if records:
                async with async_session_maker() as session:
                    await upsert_messages(session, records)
                    await session.commit()
                total_upserted += len(records)

            if (i + 1) % 10 == 0:
                logger.info(f"  [{i+1}/{min(max_conversations, len(conversations))}] msgs total: {total_fetched}")

            # Human-like delay
            await asyncio.sleep(2 + (i % 3))

    logger.info(f"DONE: messages fetched={total_fetched}, upserted={total_upserted}")
    return total_fetched, total_upserted


def _normalize_msg_type(raw_type) -> str | None:
    """Normalize a raw message type to our canonical type string."""
    if raw_type is None:
        return None
    t = str(raw_type)
    mapping = {"1": "text", "3": "image_legacy", "34": "voice", "43": "video", "47": "emoji", "49": "file"}
    return mapping.get(t, "unknown")


async def main():
    parser = argparse.ArgumentParser(description="页面浏览器抓取器")
    parser.add_argument("--entity", required=True, choices=["friends", "groups", "messages"], help="抓取实体")
    parser.add_argument("--max-employees", type=int, default=100, help="最多遍历多少员工")
    parser.add_argument("--max-conversations", type=int, default=100, help="消息模式：最多遍历多少对话")
    args = parser.parse_args()

    cookie = os.environ.get("YK_COOKIE", "")
    if not cookie:
        # Try reading from .env
        from app.config import get_settings
        cookie = get_settings().yk_cookie
    if not cookie:
        logger.error("YK_COOKIE not set in .env or environment")
        return

    logger.info(f"Starting page-based scrape: entity={args.entity}")

    if args.entity == "friends":
        await scrape_friends(cookie, args.max_employees)
    elif args.entity == "groups":
        await scrape_groups(cookie, args.max_employees)
    elif args.entity == "messages":
        await scrape_messages(cookie, args.max_conversations)


if __name__ == "__main__":
    asyncio.run(main())