xhs_crawler_claude/app/scraper.py

"""
XiaoHongShu Scraper - 核心采集逻辑
支持单线程和多线程模式
"""
import asyncio
import json
import os
import random
from datetime import datetime
from typing import List, Dict, Optional, AsyncIterator
from playwright.async_api import async_playwright, Browser, BrowserContext, Page, TimeoutError as PlaywrightTimeoutError
from app import config
from app.log_manager import LogManager


class RiskControlException(Exception):
    """风控异常"""
    def __init__(self, reason: str):
        self.reason = reason
        super().__init__(f"触发风控: {reason}")


class XiaoHongShuScraper:
    """小红书采集器"""

    def __init__(self, event_queue: asyncio.Queue, log_manager: LogManager, browser=None):
        self.event_queue = event_queue
        self.log_manager = log_manager
        self.browser = browser  # 使用传入的浏览器实例
        self.is_running = False
        self._stop_requested = False

    async def init_browser(self):
        """初始化浏览器（如果未传入浏览器实例）"""
        if not self.browser:
            self.playwright = await async_playwright().start()
            self.browser = await self.playwright.chromium.launch(
                headless=config.HEADLESS,
                args=config.BROWSER_ARGS
            )
            await self.emit_log("SUCCESS", "浏览器初始化完成")

    async def close_browser(self):
        """关闭浏览器"""
        if self.browser:
            await self.browser.close()
        if self.playwright:
            await self.playwright.stop()
        await self.emit_log("INFO", "浏览器已关闭")

    async def emit_log(self, level: str, message: str, thread_id: Optional[str] = None):
        """发送日志事件"""
        # 持久化到日志管理器
        log_entry = self.log_manager.append(level, message, thread_id)

        # 推送到 Web 前端
        await self.event_queue.put({
            "type": "log",
            **log_entry
        })

    async def emit_progress(self, **kwargs):
        """发送进度更新"""
        await self.event_queue.put({
            "type": "progress",
            **kwargs
        })

    async def emit_user_status(self, user_id: str, status: str, note_count: int = 0, thread_id: Optional[str] = None):
        """发送用户状态"""
        await self.event_queue.put({
            "type": "user_status",
            "user_id": user_id,
            "status": status,
            "note_count": note_count,
            "thread_id": thread_id
        })

    async def emit_thread_status(self, thread_id: str, status: str, current_user: Optional[str] = None):
        """发送线程状态"""
        await self.event_queue.put({
            "type": "thread_status",
            "thread_id": thread_id,
            "status": status,
            "current_user": current_user
        })

    async def handle_login(self, context: BrowserContext) -> bool:
        """
        处理登录
        返回是否已登录
        """
        page = await context.new_page()
        await page.goto(config.LOGIN_URL)

        await self.emit_log("WARNING", "请在浏览器中完成扫码登录...")
        await self.emit_log("INFO", "等待登录确认...")

        # 等待登录成功（检查 URL 变化或特定元素）
        try:
            # 等待跳转到首页或出现用户头像
            await page.wait_for_url(
                f"{config.XIAOHONGSHU_BASE_URL}/**",
                timeout=config.LOGIN_TIMEOUT * 1000
            )
            await self.emit_log("SUCCESS", "登录成功！")

            # 保存登录状态
            await context.storage_state(path=config.AUTH_STATE_FILE)
            await self.emit_log("INFO", "登录状态已保存")

            await page.close()
            return True

        except PlaywrightTimeoutError:
            await self.emit_log("ERROR", "登录超时，请重试")
            await page.close()
            return False

    async def copy_cookies(self, target_context: BrowserContext, source_context: BrowserContext):
        """复制 Cookies 到目标 Context"""
        cookies = await source_context.cookies()
        await target_context.add_cookies(cookies)

    async def check_risk_control(self, page: Page) -> tuple[bool, str]:
        """
        检测页面是否触发风控
        返回 (是否触发, 触发原因)
        """
        try:
            content = await page.content()
            for keyword in config.RISK_KEYWORDS:
                if keyword in content:
                    return True, f"检测到关键词: {keyword}"
            return False, ""
        except Exception as e:
            return False, f"检测失败: {str(e)}"

    async def human_like_scroll(self, page: Page, thread_id: Optional[str] = None):
        """
        模拟人类间歇性平滑滚动
        """
        await self.emit_log("DEBUG", f"[{thread_id}] 开始平滑滚动加载笔记...", thread_id=thread_id)

        scroll_count = 0
        last_height = 0

        while scroll_count < config.MAX_SCROLL_COUNT:
            # 获取当前页面高度
            current_height = await page.evaluate("document.body.scrollHeight")

            # 如果高度没有变化，可能到底了
            if current_height == last_height:
                await self.emit_log("DEBUG", f"[{thread_id}] 已滚动到底部", thread_id=thread_id)
                break

            last_height = current_height

            # 随机滚动距离
            scroll_step = random.randint(config.SCROLL_STEP_MIN, config.SCROLL_STEP_MAX)
            scroll_position = await page.evaluate("window.scrollY")

            # 平滑滚动
            await page.evaluate(f'''
                window.scrollTo({{
                    top: {scroll_position + scroll_step},
                    behavior: 'smooth'
                }});
            ''')

            # 随机暂停
            pause_time = random.uniform(config.SCROLL_PAUSE_MIN, config.SCROLL_PAUSE_MAX)
            await asyncio.sleep(pause_time)

            scroll_count += 1

            # 每滚动几次检查一下风控
            if scroll_count % 5 == 0:
                is_risk, reason = await self.check_risk_control(page)
                if is_risk:
                    raise RiskControlException(reason)

        await self.emit_log("DEBUG", f"[{thread_id}] 滚动完成，共滚动 {scroll_count} 次", thread_id=thread_id)

    async def extract_note_cards(self, page: Page, thread_id: Optional[str] = None) -> List[Dict]:
        """
        从用户主页提取笔记卡片
        返回 [{'url': '...', 'title': '...'}]
        """
        await self.emit_log("DEBUG", f"[{thread_id}] 提取笔记卡片...", thread_id=thread_id)

        note_urls = []

        try:
            # 尝试多种选择器
            selectors = [
                "a[href*='/explore/']",
                "section[class*='note'] a[href*='/explore/']",
                "div[class*='note-item'] a[href*='/explore/']",
            ]

            links = []
            for selector in selectors:
                try:
                    elements = await page.query_selector_all(selector)
                    if elements:
                        links = elements
                        await self.emit_log("DEBUG", f"[{thread_id}] 使用选择器: {selector}", thread_id=thread_id)
                        break
                except Exception:
                    continue

            # 提取 URL
            for link in links:
                url = await link.get_attribute('href')
                if url and '/explore/' in url:
                    # 补全 URL
                    if url.startswith('/'):
                        url = f"{config.XIAOHONGSHU_BASE_URL}{url}"

                    # 尝试获取标题
                    try:
                        title_elem = await link.query_selector("span, div")
                        title = await title_elem.inner_text() if title_elem else "无标题"
                    except Exception:
                        title = "无标题"

                    note_urls.append({
                        'url': url,
                        'title': title[:50]  # 限制标题长度
                    })

            await self.emit_log("INFO", f"[{thread_id}] 提取到 {len(note_urls)} 条笔记", thread_id=thread_id)

            return note_urls

        except Exception as e:
            await self.emit_log("ERROR", f"[{thread_id}] 提取笔记卡片失败: {str(e)}", thread_id=thread_id)
            return []

    async def extract_note_detail(self, page: Page, thread_id: Optional[str] = None) -> Optional[Dict]:
        """
        从笔记详情页提取数据
        """
        try:
            # 等待页面加载
            await asyncio.sleep(random.uniform(*config.PAGE_LOAD_DELAY))

            # 检查风控
            is_risk, reason = await self.check_risk_control(page)
            if is_risk:
                raise RiskControlException(reason)

            data = {}

            # 提取标题
            try:
                title_elem = await page.query_selector(config.SELECTOR_DETAIL_TITLE)
                data['title'] = await title_elem.inner_text() if title_elem else "无标题"
            except Exception:
                data['title'] = "无标题"

            # 提取发布时间
            try:
                time_elem = await page.query_selector(config.SELECTOR_DETAIL_PUBLISH_TIME)
                data['publish_time'] = await time_elem.inner_text() if time_elem else ""
            except Exception:
                data['publish_time'] = ""

            # 提取发布地
            try:
                loc_elem = await page.query_selector(config.SELECTOR_DETAIL_LOCATION)
                data['location'] = await loc_elem.inner_text() if loc_elem else ""
            except Exception:
                data['location'] = ""

            # 提取互动数据（尝试从页面内容解析）
            content = await page.content()

            # 评论数
            data['comment_count'] = self._extract_count(content, ['评论', 'comment'])

            # 点赞数
            data['like_count'] = self._extract_count(content, ['点赞', '赞', 'like'])

            # 收藏数
            data['collect_count'] = self._extract_count(content, ['收藏', 'star', 'collect'])

            # 笔记链接
            data['note_url'] = page.url

            return data

        except RiskControlException:
            raise
        except Exception as e:
            await self.emit_log("ERROR", f"[{thread_id}] 提取详情失败: {str(e)}", thread_id=thread_id)
            return None

    def _extract_count(self, content: str, keywords: List[str]) -> int:
        """从页面内容提取计数"""
        import re
        for keyword in keywords:
            # 匹配数字
            pattern = f"{keyword}[\\s:：]*([\\d.]+)"
            match = re.search(pattern, content)
            if match:
                count_str = match.group(1)
                # 处理万、k等单位
                if '万' in count_str:
                    return int(float(count_str.replace('万', '')) * 10000)
                elif 'k' in count_str.lower():
                    return int(float(count_str.lower().replace('k', '')) * 1000)
                else:
                    return int(count_str)
        return 0

    async def save_user_data(self, user_id: str, new_notes: List[Dict]):
        """保存用户数据（增量）"""
        file_path = os.path.join(config.DATA_DIR, f"{user_id}.json")

        # 读取现有数据
        existing_notes = []
        if os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    existing_notes = data.get('notes', [])
            except Exception as e:
                await self.emit_log("WARNING", f"读取现有数据失败: {e}")

        # 提取现有 URL
        existing_urls = {note['note_url'] for note in existing_notes}

        # 追加新笔记
        added_count = 0
        for note in new_notes:
            if note['note_url'] not in existing_urls:
                existing_notes.append(note)
                existing_urls.add(note['note_url'])
                added_count += 1

        # 保存
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump({
                "user_id": user_id,
                "last_updated": datetime.now().isoformat(),
                "total_notes": len(existing_notes),
                "notes": existing_notes
            }, f, ensure_ascii=False, indent=2)

        return added_count

    async def process_user(self, context: BrowserContext, user_id: str, thread_id: str = "Main") -> int:
        """
        处理单个用户
        返回新增笔记数
        """
        await self.emit_user_status(user_id, "processing", thread_id=thread_id)
        await self.emit_thread_status(thread_id, "busy", user_id)

        try:
            await self.emit_log("INFO", f"[{thread_id}] 开始处理用户: {user_id}", thread_id=thread_id)

            # 创建新页面
            page = await context.new_page()

            # 访问用户主页
            url = config.USER_PROFILE_URL_TEMPLATE.format(user_id=user_id)
            await self.emit_log("INFO", f"[{thread_id}] 访问用户主页: {url}", thread_id=thread_id)

            try:
                await page.goto(url, timeout=30000)
            except PlaywrightTimeoutError:
                await self.emit_log("ERROR", f"[{thread_id}] 访问用户主页超时: {user_id}", thread_id=thread_id)
                await page.close()
                await self.emit_user_status(user_id, "failed", thread_id=thread_id)
                await self.emit_thread_status(thread_id, "idle", None)
                return 0

            # 检查风控
            is_risk, reason = await self.check_risk_control(page)
            if is_risk:
                await self.emit_log("ERROR", f"[{thread_id}] 访问用户主页触发风控: {reason}", thread_id=thread_id)
                await page.close()
                await self.emit_user_status(user_id, "failed", thread_id=thread_id)
                await self.emit_thread_status(thread_id, "idle", None)
                raise RiskControlException(reason)

            await asyncio.sleep(random.uniform(*config.PAGE_LOAD_DELAY))

            # 滚动加载所有笔记
            await self.human_like_scroll(page, thread_id)

            # 提取笔记卡片
            note_cards = await self.extract_note_cards(page, thread_id)

            if not note_cards:
                await self.emit_log("WARNING", f"[{thread_id}] 未找到笔记，可能页面结构变化", thread_id=thread_id)
                await page.close()
                await self.emit_user_status(user_id, "completed", 0, thread_id=thread_id)
                await self.emit_thread_status(thread_id, "idle", None)
                return 0

            # 采集每条笔记
            collected_notes = []
            for i, card in enumerate(note_cards):
                if self._stop_requested:
                    await self.emit_log("WARNING", f"[{thread_id}] 收到停止信号", thread_id=thread_id)
                    break

                await self.emit_log("INFO", f"[{thread_id}] 采集笔记 {i+1}/{len(note_cards)}: {card['title']}", thread_id=thread_id)

                try:
                    # 访问详情页
                    await page.goto(card['url'], timeout=30000)

                    # 模拟阅读停留
                    await asyncio.sleep(random.uniform(*config.NOTE_READ_DELAY))

                    # 提取数据
                    note_data = await self.extract_note_detail(page, thread_id)

                    if note_data:
                        collected_notes.append(note_data)
                        await self.emit_log("SUCCESS", f"[{thread_id}] 成功采集: {note_data['title'][:30]}", thread_id=thread_id)

                    # 返回列表页
                    await page.go_back(url, wait_until="networkidle")

                    # 间隔
                    await asyncio.sleep(random.uniform(*config.NOTE_INTERVAL_DELAY))

                except RiskControlException:
                    await self.emit_log("ERROR", f"[{thread_id}] 采集笔记触发风控", thread_id=thread_id)
                    raise
                except Exception as e:
                    await self.emit_log("WARNING", f"[{thread_id}] 采集笔记失败: {str(e)}", thread_id=thread_id)
                    continue

            await page.close()

            # 保存数据
            if collected_notes:
                added_count = await self.save_user_data(user_id, collected_notes)
                await self.emit_log("SUCCESS", f"[{thread_id}] 用户 {user_id} 采集完成，新增 {added_count} 条笔记", thread_id=thread_id)
                await self.emit_user_status(user_id, "completed", len(collected_notes), thread_id=thread_id)
            else:
                await self.emit_log("INFO", f"[{thread_id}] 用户 {user_id} 无新笔记", thread_id=thread_id)
                await self.emit_user_status(user_id, "completed", 0, thread_id=thread_id)

            await self.emit_thread_status(thread_id, "idle", None)
            return len(collected_notes)

        except RiskControlException:
            await self.emit_user_status(user_id, "failed", thread_id=thread_id)
            await self.emit_thread_status(thread_id, "error", None)
            raise
        except Exception as e:
            await self.emit_log("ERROR", f"[{thread_id}] 处理用户 {user_id} 失败: {str(e)}", thread_id=thread_id)
            await self.emit_user_status(user_id, "failed", thread_id=thread_id)
            await self.emit_thread_status(thread_id, "error", None)
            return 0

    async def run_single_thread(self, user_ids: List[str]):
        """单线程模式"""
        await self.emit_log("INFO", "启动单线程采集模式")

        # 创建 Context 并登录
        
        # 创建 Context（优先加载已保存的 storage_state）
        if os.path.exists(config.AUTH_STATE_FILE):
            await self.emit_log("INFO", "加载已保存的登录状态 (storage_state)")
            context = await self.browser.new_context(
                storage_state=config.AUTH_STATE_FILE,
                viewport={'width': config.VIEWPORT_WIDTH, 'height': config.VIEWPORT_HEIGHT},
                user_agent=config.USER_AGENT
            )
        else:
            context = await self.browser.new_context(
                viewport={'width': config.VIEWPORT_WIDTH, 'height': config.VIEWPORT_HEIGHT},
                user_agent=config.USER_AGENT
            )
            # 需要登录
            login_success = await self.handle_login(context)
            if not login_success:
                await context.close()
                return

        self.is_running = True

        for i, user_id in enumerate(user_ids):
            if self._stop_requested:
                await self.emit_log("WARNING", "收到停止请求")
                break

            await self.emit_log("INFO", f"处理用户 {i+1}/{len(user_ids)}: {user_id}")

            try:
                await self.process_user(context, user_id, thread_id="Main")

                # 用户切换延迟
                if i < len(user_ids) - 1:
                    delay = random.uniform(*config.USER_SWITCH_DELAY)
                    await self.emit_log("INFO", f"等待 {delay:.1f} 秒后处理下一个用户...")
                    await asyncio.sleep(delay)

            except RiskControlException as e:
                await self.emit_log("ERROR", f"触发风控: {e.reason}，终止采集")
                break
            except Exception as e:
                await self.emit_log("ERROR", f"处理用户 {user_id} 失败: {str(e)}")
                continue

        await context.close()
        self.is_running = False
        await self.emit_log("SUCCESS", "单线程采集完成")

    async def worker(self, context: BrowserContext, task_queue: asyncio.Queue, thread_id: str):
        """
        Worker 处理任务（多线程模式）
        """
        # 线程启动时错开
        await asyncio.sleep(int(thread_id.split('-')[1]) * config.THREAD_OFFSET)

        await self.emit_thread_status(thread_id, "idle", None)

        while True:
            if self._stop_requested:
                break

            try:
                # 获取任务
                user_id = await asyncio.wait_for(task_queue.get(), timeout=1.0)
                await self.emit_thread_status(thread_id, "busy", user_id)

                try:
                    # 线程间最小间隔
                    await asyncio.sleep(random.uniform(
                        config.THREAD_MIN_DELAY,
                        config.THREAD_MIN_DELAY + 3
                    ))

                    await self.process_user(context, user_id, thread_id)

                except RiskControlException as e:
                    await self.emit_log("ERROR", f"[{thread_id}] 触发风控: {e.reason}，停止该线程", thread_id=thread_id)
                    await self.emit_thread_status(thread_id, "error", None)
                    break
                except Exception as e:
                    await self.emit_log("ERROR", f"[{thread_id}] 处理失败: {str(e)}", thread_id=thread_id)
                finally:
                    task_queue.task_done()
                    await self.emit_thread_status(thread_id, "idle", None)

            except asyncio.TimeoutError:
                continue
            except Exception as e:
                await self.emit_log("ERROR", f"[{thread_id}] Worker 错误: {str(e)}", thread_id=thread_id)
                break

    async def run_multi_thread(self, user_ids: List[str], max_threads: int = 3):
        """多线程模式"""
        await self.emit_log("INFO", f"启动多线程采集模式，线程数: {max_threads}")


        # 创建多个 Context（同一账号并发会增加风控风险，请合理控制线程数）
        contexts = []

        # 确保已有登录态；若没有，先用第一个 context 登录并保存
        if not os.path.exists(config.AUTH_STATE_FILE):
            first = await self.browser.new_context(
                viewport={'width': config.VIEWPORT_WIDTH, 'height': config.VIEWPORT_HEIGHT},
                user_agent=config.USER_AGENT
            )
            login_success = await self.handle_login(first)
            if not login_success:
                await first.close()
                return
            await first.close()

        # 基于 storage_state 为每个线程创建独立 context
        for i in range(max_threads):
            ctx = await self.browser.new_context(
                storage_state=config.AUTH_STATE_FILE,
                viewport={'width': config.VIEWPORT_WIDTH, 'height': config.VIEWPORT_HEIGHT},
                user_agent=config.USER_AGENT
            )
            contexts.append(ctx)
            await self.emit_log("INFO", f"Thread-{i+1} context ready")

        # 创建任务队列
        task_queue = asyncio.Queue()
        for user_id in user_ids:
            await task_queue.put(user_id)

        self.is_running = True

        # 创建 workers
        workers = [
            asyncio.create_task(
                self.worker(contexts[i], task_queue, f"Thread-{i+1}")
            )
            for i in range(max_threads)
        ]

        # 等待完成
        await task_queue.join()

        # 取消 workers
        for worker in workers:
            worker.cancel()

        # 关闭 contexts
        for ctx in contexts:
            await ctx.close()

        self.is_running = False
        await self.emit_log("SUCCESS", "多线程采集完成")

    def stop(self):
        """请求停止采集"""
        self._stop_requested = True
        self.is_running = False

    async def run(self, user_ids: List[str], mode: str = "single", max_threads: int = 3):
        """
        运行采集
        :param user_ids: 用户 ID 列表
        :param mode: single / multi
        :param max_threads: 多线程模式下的线程数
        """
        self._stop_requested = False

        await self.emit_log("INFO", f"开始采集，共 {len(user_ids)} 个用户")

        await self.init_browser()

        try:
            if mode == "single":
                await self.run_single_thread(user_ids)
            else:
                await self.run_multi_thread(user_ids, max_threads)
        finally:
            await self.close_browser()

        await self.emit_log("SUCCESS", "所有采集任务完成")