screenshotScript.py

#!/usr/bin/env python3
"""
独立的网页截图工具
支持命令行参数：
- -u/--url: 指定单个URL进行截图
- -f/--file: 指定包含URL列表的文件进行批量截图
截图保存到当前目录下的img文件夹，并生成HTML图片链接页面
"""

import asyncio
import argparse
import os
import sys
import logging
import datetime
import json
from typing import List, Optional, Dict, Tuple
from pathlib import Path

from playwright.async_api import async_playwright

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class ScreenshotTool:
    """截图工具类"""
    
    def __init__(self, concurrency: int = 3):
        """
        初始化截图工具
        
        Args:
            concurrency: 并发截图数
        """
        self.concurrency = concurrency
        self.current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.output_dir = Path(f"img_{self.current_time}")
        
        # 创建输出目录
        self.output_dir.mkdir(exist_ok=True)
        
        # 存储URL到文件名的映射
        self.url_to_filename = {}
        # 存储文件名到URL的映射
        self.filename_to_url = {}
    
    async def capture_screenshot(self, url: str, page) -> Tuple[bool, str]:
        """
        捕获单个URL的截图
        
        Args:
            url: 目标URL
            page: Playwright Page对象
        
        Returns:
            (是否成功, 页面标题)
        """
        try:
            # 清理URL，用于文件名
            filename = self._sanitize_filename(url)
            output_path = self.output_dir / f"{filename}.png"
            
            # 保存URL和文件名的映射关系
            self.url_to_filename[url] = filename
            self.filename_to_url[filename] = url
            
            # 保存映射到文件（用于故障恢复）
            self._save_url_mapping()
            
            # 尝试加载页面
            try:
                await page.goto(url, timeout=30000, wait_until='networkidle')
                
                # 获取页面标题
                title = await page.title()
                if not title or title.strip() == '':
                    title = url  # 如果标题为空，使用URL作为标题
                else:
                    # 清理标题中的换行符和多余空格
                    title = title.replace('\n', ' ').replace('\r', '').strip()
                
                logger.info(f"获取到页面标题: {title[:50]}...")
                
            except Exception as e:
                logger.warning(f"页面加载异常但尝试截图: {url}, 错误: {str(e)[:50]}")
                title = url  # 加载异常时使用URL作为标题
            
            # 截图
            await page.screenshot(
                path=str(output_path),
                type='png',
                full_page=False
            )
            
            logger.info(f"截图成功: {url} -> {output_path}")
            return True, title
            
        except asyncio.TimeoutError:
            logger.error(f"截图超时: {url}")
            return False, url
        except Exception as e:
            logger.error(f"截图失败: {url}, 错误: {str(e)[:100]}")
            return False, url
    
    def _save_url_mapping(self):
        """保存URL映射关系到文件"""
        mapping_file = self.output_dir / "url_mapping.json"
        try:
            with open(mapping_file, 'w', encoding='utf-8') as f:
                json.dump({
                    'url_to_filename': self.url_to_filename,
                    'filename_to_url': self.filename_to_url
                }, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.warning(f"保存URL映射关系失败: {e}")
    
    def _load_url_mapping(self):
        """从文件加载URL映射关系"""
        mapping_file = self.output_dir / "url_mapping.json"
        if mapping_file.exists():
            try:
                with open(mapping_file, 'r', encoding='utf-8') as f:
                    mapping = json.load(f)
                    self.url_to_filename = mapping.get('url_to_filename', {})
                    self.filename_to_url = mapping.get('filename_to_url', {})
            except Exception as e:
                logger.warning(f"加载URL映射关系失败: {e}")
    
    async def _capture_with_semaphore(self, url: str, context, semaphore: asyncio.Semaphore) -> tuple:
        """
        使用信号量控制并发的截图任务
        
        Args:
            url: 目标URL
            context: Playwright BrowserContext
            semaphore: 并发控制信号量
        
        Returns:
            (url, 是否成功, 页面标题) 元组
        """
        async with semaphore:
            page = await context.new_page()
            try:
                success, title = await self.capture_screenshot(url, page)
                return (url, success, title)
            finally:
                await page.close()
    
    async def capture_single_url(self, url: str) -> Tuple[bool, str]:
        """
        截图单个URL
        
        Args:
            url: 目标URL
        
        Returns:
            (是否成功, 页面标题)
        """
        async with async_playwright() as p:
            browser = await p.chromium.launch(
                headless=True,
                args=[
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--disable-dev-shm-usage',
                    '--disable-gpu'
                ]
            )
            
            try:
                context = await browser.new_context(
                    viewport={'width': 1920, 'height': 1080},
                    ignore_https_errors=True,
                    user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
                )
                
                page = await context.new_page()
                success, title = await self.capture_screenshot(url, page)
                await page.close()
                await context.close()
                
                return success, title
                
            finally:
                await browser.close()
    
    async def capture_urls_from_file(self, file_path: str) -> Dict:
        """
        从文件读取URL列表进行批量截图
        
        Args:
            file_path: 包含URL列表的文件路径
        
        Returns:
            统计信息字典和标题映射
        """
        urls = self._read_urls_from_file(file_path)
        if not urls:
            logger.error(f"文件 {file_path} 中没有有效的URL")
            return {'total': 0, 'successful': 0, 'failed': 0}
        
        return await self.capture_urls(urls)
    
    async def capture_urls(self, urls: List[str]) -> Dict:
        """
        批量截图URL列表
        
        Args:
            urls: URL列表
        
        Returns:
            统计信息字典
        """
        if not urls:
            return {'total': 0, 'successful': 0, 'failed': 0}
        
        total = len(urls)
        successful = 0
        failed = 0
        
        # 存储URL和标题的映射
        url_title_map = {}
        
        logger.info(f"开始批量截图 - URL数: {total}, 并发数: {self.concurrency}")
        
        async with async_playwright() as p:
            browser = await p.chromium.launch(
                headless=True,
                args=[
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--disable-dev-shm-usage',
                    '--disable-gpu'
                ]
            )
            
            try:
                context = await browser.new_context(
                    viewport={'width': 1920, 'height': 1080},
                    ignore_https_errors=True,
                    user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
                )
                
                # 使用信号量控制并发
                semaphore = asyncio.Semaphore(self.concurrency)
                
                # 创建所有任务
                tasks = [
                    self._capture_with_semaphore(url, context, semaphore)
                    for url in urls
                ]
                
                # 执行所有任务
                for coro in asyncio.as_completed(tasks):
                    url, success, title = await coro
                    
                    # 存储标题信息
                    url_title_map[url] = title
                    
                    if success:
                        successful += 1
                    else:
                        failed += 1
                    
                    # 每10个URL打印一次进度
                    if (successful + failed) % 10 == 0:
                        logger.info(f"进度: {successful + failed}/{total}")
                
                await context.close()
                
            finally:
                await browser.close()
        
        logger.info(f"截图完成 - 总数: {total}, 成功: {successful}, 失败: {failed}")
        
        # 生成HTML图片链接页面
        if successful > 0:
            self._generate_html_gallery(url_title_map)
        
        return {'total': total, 'successful': successful, 'failed': failed}
    
    def _read_urls_from_file(self, file_path: str) -> List[str]:
        """
        从文件读取URL列表
        
        Args:
            file_path: 文件路径
        
        Returns:
            URL列表
        """
        if not os.path.exists(file_path):
            logger.error(f"文件不存在: {file_path}")
            return []
        
        urls = []
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and self._is_valid_url(line):
                        urls.append(line)
        except Exception as e:
            logger.error(f"读取文件失败: {file_path}, 错误: {e}")
            return []
        
        return urls
    
    def _is_valid_url(self, url: str) -> bool:
        """
        检查URL是否有效
        
        Args:
            url: URL字符串
        
        Returns:
            是否有效
        """
        return url.startswith(('http://', 'https://'))
    
    def _sanitize_filename(self, url: str) -> str:
        """
        清理URL，生成安全的文件名（改进版）
        
        Args:
            url: URL字符串
        
        Returns:
            安全的文件名
        """
        import base64
        import hashlib
        
        # 使用base64编码作为文件名，这是可逆的
        # base64.urlsafe_b64encode会将+和/替换为-和_，已经是安全的
        encoded = base64.urlsafe_b64encode(url.encode('utf-8')).decode('ascii')
        
        # 移除base64的填充字符=，用_EQ_代替以便恢复
        filename = encoded.rstrip('=').replace('=', '_EQ_')
        
        # 如果文件名仍然太长，使用哈希值（但这样不可逆）
        if len(filename) > 200:
            # 使用URL的MD5哈希作为文件名（16进制，32字符）
            filename = hashlib.md5(url.encode('utf-8')).hexdigest()
            logger.warning(f"URL过长，使用MD5哈希作为文件名: {url[:50]}... -> {filename}")
        
        return filename
    
    def _filename_to_url(self, filename: str) -> str:
        """
        从文件名还原URL（修复版）
        
        Args:
            filename: 文件名（不含扩展名）
        
        Returns:
            原始URL
        """
        # 先尝试从映射中查找
        if filename in self.filename_to_url:
            return self.filename_to_url[filename]
        
        # 如果没有映射，尝试从文件加载
        self._load_url_mapping()
        if filename in self.filename_to_url:
            return self.filename_to_url[filename]
        
        # 如果仍然没有，尝试解码base64
        try:
            import base64
            
            # 恢复_EQ_为=
            encoded = filename.replace('_EQ_', '=')
            
            # 添加base64填充字符
            padding = 4 - len(encoded) % 4
            if padding != 4:
                encoded += '=' * padding
            
            # 解码
            decoded_bytes = base64.urlsafe_b64decode(encoded)
            url = decoded_bytes.decode('utf-8')
            return url
        except Exception as e:
            # 如果解码失败，尝试旧的还原方法（不完美）
            logger.warning(f"无法解码文件名 {filename}: {e}")
            
            # 尝试旧的还原方法
            if filename.startswith('http__'):
                url = 'http://' + filename[6:]
            elif filename.startswith('https__'):
                url = 'https://' + filename[7:]
            else:
                # 可能是MD5哈希，无法还原
                return f"[无法还原原始URL: {filename}]"
            
            # 简单替换下划线为斜杠（不完美，但有bug）
            # 这里应该只替换路径部分的下划线，但简单实现中难以区分
            url = url.replace('_', '/')
            return url
    
    def _generate_html_gallery(self, url_title_map: Dict[str, str] = None) -> None:
        """
        生成HTML图片链接页面，每行显示4个图片
        
        Args:
            url_title_map: URL到标题的映射字典
        """
        # 获取当前时间，精确到秒
        html_filename = f"screenshots_{self.current_time}.html"
        
        # 获取所有PNG文件
        png_files = list(self.output_dir.glob("*.png"))
        if not png_files:
            logger.warning("没有找到PNG文件，跳过HTML生成")
            return
        
        # 按文件名排序
        png_files.sort()
        
        # 生成HTML内容
        html_content = self._build_html_content(png_files, url_title_map)
        
        # 保存HTML文件
        with open(html_filename, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        logger.info(f"HTML图片链接页面已生成: {html_filename}")
    
    def _build_html_content(self, png_files: list, url_title_map: Dict[str, str] = None) -> str:
        """
        构建HTML内容
        
        Args:
            png_files: PNG文件列表
            url_title_map: URL到标题的映射字典
        
        Returns:
            HTML内容字符串
        """
        html = '''<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>网页截图预览</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            margin: 20px;
            background-color: #f5f5f5;
        }
        .container {
            max-width: 1400px;
            margin: 0 auto;
            background: white;
            padding: 20px;
            border-radius: 8px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        }
        h1 {
            text-align: center;
            color: #333;
            margin-bottom: 30px;
        }
        .gallery {
            display: grid;
            grid-template-columns: repeat(4, 1fr);
            gap: 20px;
        }
        .screenshot-item {
            text-align: center;
            border: 1px solid #ddd;
            border-radius: 8px;
            padding: 15px;
            background: #fafafa;
            transition: transform 0.2s;
            display: flex;
            flex-direction: column;
            height: 320px;
            overflow: hidden;
        }
        .screenshot-item:hover {
            transform: translateY(-3px);
            box-shadow: 0 6px 12px rgba(0,0,0,0.15);
        }
        .screenshot-item img {
            width: 100%;
            height: 200px;
            object-fit: cover;
            border-radius: 4px;
            cursor: pointer;
            margin-bottom: 10px;
        }
        .screenshot-item .title {
            font-weight: bold;
            font-size: 14px;
            color: #333;
            margin-bottom: 5px;
            height: 40px;
            overflow: hidden;
            display: -webkit-box;
            -webkit-line-clamp: 2;
            -webkit-box-orient: vertical;
            line-clamp: 2;
        }
        .screenshot-item .url {
            font-size: 12px;
            color: #666;
            word-break: break-all;
            max-height: 40px;
            overflow: hidden;
            display: -webkit-box;
            -webkit-line-clamp: 2;
            -webkit-box-orient: vertical;
            line-clamp: 2;
        }
        .screenshot-item .url a {
            color: #0066cc;
            text-decoration: none;
        }
        .screenshot-item .url a:hover {
            text-decoration: underline;
        }
        .modal {
            display: none;
            position: fixed;
            z-index: 1000;
            left: 0;
            top: 0;
            width: 100%;
            height: 100%;
            background-color: rgba(0,0,0,0.9);
        }
        .modal-content {
            margin: auto;
            display: block;
            max-width: 90%;
            max-height: 90%;
            margin-top: 5%;
        }
        .close {
            position: absolute;
            top: 15px;
            right: 35px;
            color: #fff;
            font-size: 40px;
            font-weight: bold;
            cursor: pointer;
        }
        .stats {
            text-align: center;
            margin-bottom: 20px;
            color: #666;
            padding: 10px;
            background: #f0f0f0;
            border-radius: 6px;
        }
        @media (max-width: 1200px) {
            .gallery {
                grid-template-columns: repeat(3, 1fr);
            }
        }
        @media (max-width: 900px) {
            .gallery {
                grid-template-columns: repeat(2, 1fr);
            }
        }
        @media (max-width: 600px) {
            .gallery {
                grid-template-columns: 1fr;
            }
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>📸 网页截图预览</h1>
        <div class="stats">
            共 <strong>''' + str(len(png_files)) + '''</strong> 张截图，生成时间: ''' + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + '''
        </div>
        <div class="gallery">
'''
        
        # 添加图片项
        for png_file in png_files:
            # 获取文件名（不含扩展名）
            filename_without_ext = png_file.stem
            
            # 从映射中获取URL
            url = self._filename_to_url(filename_without_ext)
            
            # 获取页面标题，如果没有则使用URL
            title = url_title_map.get(url, url) if url_title_map else url
            # 限制标题长度
            if len(title) > 80:
                display_title = title[:77] + "..."
            else:
                display_title = title
            
            # 限制URL显示长度
            display_url = url
            if len(display_url) > 60:
                display_url = display_url[:57] + "..."
            
            html += f'''            <div class="screenshot-item">
                <img src="{self.output_dir}/{png_file.name}" alt="{title}" onclick="openModal(this)">
                <div class="title" title="{title}">{display_title}</div>
                <div class="url" title="{url}"><a href="{url}" target="_blank">{display_url}</a></div>
            </div>
'''
        
        html += '''        </div>
    </div>
    
    <!-- 模态框 -->
    <div id="modal" class="modal">
        <span class="close" onclick="closeModal()">&times;</span>
        <img class="modal-content" id="modal-img">
    </div>
    
    <script>
        function openModal(img) {
            document.getElementById('modal').style.display = 'block';
            document.getElementById('modal-img').src = img.src;
            document.body.style.overflow = 'hidden'; // 防止背景滚动
        }
        
        function closeModal() {
            document.getElementById('modal').style.display = 'none';
            document.body.style.overflow = 'auto'; // 恢复背景滚动
        }
        
        // 点击模态框背景关闭
        document.getElementById('modal').onclick = function(event) {
            if (event.target === this) {
                closeModal();
            }
        }
        
        // ESC键关闭
        document.addEventListener('keydown', function(event) {
            if (event.key === 'Escape') {
                closeModal();
            }
        });
    </script>
</body>
</html>'''
        
        return html


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='网页截图工具')
    parser.add_argument('-u', '--url', type=str, help='单个URL进行截图')
    parser.add_argument('-f', '--file', type=str, help='包含URL列表的文件路径')
    parser.add_argument('-c', '--concurrency', type=int, default=3, help='并发数（默认3）')
    
    args = parser.parse_args()
    
    # 检查参数
    if not args.url and not args.file:
        parser.print_help()
        print("\n错误: 必须指定 -u/--url 或 -f/--file 参数")
        sys.exit(1)
    
    if args.url and args.file:
        print("错误: 不能同时使用 -u 和 -f 参数")
        sys.exit(1)
    
    # 创建截图工具
    tool = ScreenshotTool(concurrency=args.concurrency)
    
    # 执行截图
    try:
        if args.url:
            # 截图单个URL
            success, title = asyncio.run(tool.capture_single_url(args.url))
            if success:
                print(f"截图成功: {args.url}")
                print(f"页面标题: {title}")
                # 单个URL也生成HTML
                url_title_map = {args.url: title}
                tool._generate_html_gallery(url_title_map)
            else:
                print(f"截图失败: {args.url}")
                sys.exit(1)
        else:
            # 批量截图
            result = asyncio.run(tool.capture_urls_from_file(args.file))
            print(f"批量截图完成: 总数={result['total']}, 成功={result['successful']}, 失败={result['failed']}")
            
    except KeyboardInterrupt:
        print("\n用户中断操作")
        sys.exit(1)
    except Exception as e:
        logger.error(f"程序执行异常: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()