kb-skills/kb-zhihu-content/zhihu_content.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知乎内容获取工具

支持：
- zvideo: 知乎视频下载
- 从页面提取视频播放地址
"""

import argparse
import asyncio
import json
import os
import re
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any, List


class ZhihuContent:
    """知乎内容获取器"""

    def __init__(self, output_dir: str = "./output", cookie: str = None):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.cookie = cookie

    def _create_cookie_file(self) -> str:
        """创建 Netscape 格式的 Cookie 文件"""
        if not self.cookie:
            return None

        cookies = {}
        for item in self.cookie.split(';'):
            item = item.strip()
            if '=' in item:
                key, value = item.split('=', 1)
                cookies[key.strip()] = value.strip()

        cookie_content = """# Netscape HTTP Cookie File
# Generated by zhihu_content.py

"""

        for name, value in cookies.items():
            cookie_content += f".zhihu.com\tTRUE\t/\tFALSE\t0\t{name}\t{value}\n"

        fd, path = tempfile.mkstemp(suffix='.txt')
        with os.fdopen(fd, 'w') as f:
            f.write(cookie_content)

        return path

    async def download_video(self, video_id: str, quality: int = 720) -> Dict[str, Any]:
        """下载知乎视频（使用 yt-dlp）"""
        result = {"success": False, "video_id": video_id, "message": "", "filepath": None, "metadata": {}}

        try:
            print(f"获取知乎视频信息: {video_id}")

            cookie_file = self._create_cookie_file()
            url = f"https://www.zhihu.com/zvideo/{video_id}"

            # 先获取视频标题
            try:
                if cookie_file:
                    cmd = ["yt-dlp", "--cookies", cookie_file, "--get-title", url]
                else:
                    cmd = ["yt-dlp", "--get-title", url]

                proc = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE
                )
                stdout, stderr = await proc.communicate()

                title = stdout.decode('utf-8').strip()
                if not title or 'error' in title.lower():
                    title = f"zhihu_video_{video_id}"
            except Exception as e:
                title = f"zhihu_video_{video_id}"

            # 构建 yt-dlp 命令
            safe_title = "".join(c for c in title if c.isalnum() or c in (' ', '-', '_')).strip()[:50]

            print(f"下载视频: {title}")
            output_template = str(self.output_dir / f"{safe_title}.%(ext)s")

            # 构建下载命令
            cmd = [
                "yt-dlp",
                "-f", "bestvideo+bestaudio/best",
                "-o", output_template,
                "--no-download-archive",
                url
            ]

            if cookie_file:
                cmd.extend(["--cookies", cookie_file])

            print(f"执行命令: {' '.join(cmd[:4])}...")

            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await proc.communicate()

            # 查找下载的文件
            downloaded_files = list(self.output_dir.glob(f"{safe_title}.*"))
            if downloaded_files:
                video_file = downloaded_files[0]
                result["success"] = True
                result["filepath"] = str(video_file)
                result["message"] = f"下载完成: {video_file.name}"
                result["metadata"] = {"title": title, "video_id": video_id}
                print(f"✅ {result['message']}")
            else:
                result["message"] = "未找到下载的文件"
                print(f"❌ {result['message']}")

            # 清理临时 Cookie 文件
            if cookie_file and os.path.exists(cookie_file):
                os.remove(cookie_file)

        except Exception as e:
            result["message"] = f"错误: {str(e)}"
            print(f"❌ {result['message']}")

        return result

    async def extract_subtitle(self, video_id: str) -> Dict[str, Any]:
        """
        提取字幕
        
        知乎视频通常没有官方字幕，需要使用 Whisper 进行语音识别。
        此函数返回 has_subtitle=False，让工作流继续执行 Whisper 步骤。
        """
        result = {"success": True, "video_id": video_id, "has_subtitle": False, "message": "知乎视频无官方字幕，将使用 Whisper 进行语音识别"}

        return result


async def main():
    parser = argparse.ArgumentParser(description='知乎内容获取工具')
    parser.add_argument('--id', '-i', required=True, help='知乎视频ID')
    parser.add_argument('--type', '-t', required=True, choices=['video', 'subtitle'], help='内容类型')
    parser.add_argument('--output', '-o', default='./output', help='输出目录')
    parser.add_argument('--quality', '-q', type=int, default=720, help='视频质量')

    args = parser.parse_args()

    cookie = os.environ.get('ZHIHU_COOKIE')
    downloader = ZhihuContent(output_dir=args.output, cookie=cookie)

    if args.type == 'video':
        result = await downloader.download_video(args.id, quality=args.quality)
    else:
        result = await downloader.extract_subtitle(args.id)

    # 统一输出格式
    if result.get("success"):
        output = {
            "success": True,
            "type": args.type,
            "has_subtitle": result.get("has_subtitle", False),
            "filename": Path(result.get('filepath', '')).name if result.get('filepath') else None,
            "filepath": str(Path(result.get('filepath', '')).resolve()) if result.get('filepath') else None,
            "message": result.get('message', '处理成功')
        }
    else:
        output = {
            "success": False,
            "error": result.get('message', '处理失败'),
            "detail": f"ID: {args.id}, 类型: {args.type}"
        }

    print(json.dumps(output, ensure_ascii=False, indent=2))
    return result


if __name__ == '__main__':
    try:
        result = asyncio.run(main())
        sys.exit(0 if result.get("success") else 1)
    except KeyboardInterrupt:
        print(json.dumps({"success": False, "error": "用户取消"}, ensure_ascii=False, indent=2))
        sys.exit(1)
    except Exception as e:
        print(json.dumps({"success": False, "error": str(e)}, ensure_ascii=False, indent=2))
        sys.exit(1)