Spaces:

CatPtain
/

page_shot

Paused

File size: 13,505 Bytes

const express = require('express');
const puppeteer = require('puppeteer');
const cors = require('cors');
const helmet = require('helmet');
const rateLimit = require('express-rate-limit');

const app = express();
const PORT = process.env.PORT || 7860;
const API_KEY = process.env.API_KEY; // 从环境变量获取API密钥

// 安全中间件
app.use(helmet({
    contentSecurityPolicy: {
      directives: {
        defaultSrc: ["'self'"],
        scriptSrc: ["'self'", "'unsafe-inline'"],
        styleSrc: ["'self'", "'unsafe-inline'"],
        imgSrc: ["'self'", "data:", "blob:"],
        connectSrc: ["'self'", "https:"]  // 添加 https: 允许连接外部网站
      }
    }
  }));
app.use(cors());
app.use(express.json());

// 速率限制
const limiter = rateLimit({
  windowMs: 15 * 60 * 1000, // 15分钟
  max: 100, // 限制每个IP 100次请求
  message: {
    error: 'Too many requests',
    message: 'Rate limit exceeded. Please try again later.'
  }
});
app.use(limiter);

// 队列管理
let activeRequests = 0;
const MAX_CONCURRENT = 3;
const requestQueue = [];

// CPU 监控
function getCpuUsage() {
  const usage = process.cpuUsage();
  const total = usage.user + usage.system;
  return Math.round((total / 1000000) % 100); // 简化的CPU使用率
}

// API鉴权中间件
function authenticate(req, res, next) {
  // 如果没有设置API_KEY环境变量，跳过鉴权（开发模式）
  if (!API_KEY) {
    return next();
  }

  const authHeader = req.headers.authorization;
  const apiKey = req.headers['x-api-key'] || (authHeader && authHeader.startsWith('Bearer ') ? authHeader.slice(7) : null);

  if (!apiKey || apiKey !== API_KEY) {
    return res.status(401).json({
      error: 'Unauthorized',
      message: 'Valid API key required. Please provide API key in Authorization header or x-api-key header.',
      hint: 'Use "Authorization: Bearer YOUR_API_KEY" or "x-api-key: YOUR_API_KEY"'
    });
  }

  next();
}

// 健康检查
app.get('/', (req, res) => {
  res.json({
    message: 'Page Screenshot API - Hugging Face Spaces',
    version: '1.4.0',
    status: 'running',
    platform: 'HuggingFace Spaces',
    authentication: {
      type: API_KEY ? 'API Key Required' : 'Open Access (No API Key Set)',
      required: !!API_KEY,
      note: API_KEY ? 'API key required for screenshot endpoint' : 'Set API_KEY environment variable to enable authentication'
    },
    endpoints: {
      screenshot: 'POST /screenshot (requires API key if configured)',
      demo: 'GET /demo',
      health: 'GET /',
      status: 'GET /status'
    },
    license: 'Proprietary - Commercial use requires license'
  });
});

// 状态检查
app.get('/status', (req, res) => {
  const cpuUsage = getCpuUsage();
  
  res.json({
    status: 'running',
    timestamp: new Date().toISOString(),
    system: {
      cpuUsage: `${cpuUsage}%`,
      memoryUsage: '', // HF Spaces 不提供详细内存信息
      uptime: `${Math.floor(process.uptime())}s`
    },
    queue: {
      activeRequests,
      queuedRequests: requestQueue.length,
      maxConcurrent: MAX_CONCURRENT
    },
    authentication: {
      enabled: !!API_KEY,
      type: API_KEY ? 'Environment Variable API Key' : 'Disabled',
      note: API_KEY ? 'API key authentication is active' : 'Set API_KEY environment variable to enable authentication'
    }
  });
});

// 演示页面
app.get('/demo', (req, res) => {
  res.send(`

<!DOCTYPE html>

<html>

<head>

    <title>Page Screenshot API Demo</title>

    <style>

        body { font-family: -apple-system, BlinkMacSystemFont, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }

        .form-group { margin: 15px 0; }

        input, button { padding: 12px; margin: 5px; border: 1px solid #ddd; border-radius: 6px; }

        input[type="text"] { width: 400px; }

        button { background: #007bff; color: white; border: none; cursor: pointer; }

        button:hover { background: #0056b3; }

        #result { margin-top: 20px; }

        img { max-width: 100%; border: 1px solid #ccc; border-radius: 6px; }

        .loading { color: #666; font-style: italic; }

        .error { color: #dc3545; }

        .success { color: #28a745; }

        .auth-section { background: #f8f9fa; padding: 15px; border-radius: 6px; margin-bottom: 20px; }

        .auth-info { color: #6c757d; font-size: 14px; margin-top: 10px; }

    </style>

</head>

<body>

    <h1>📸 Page Screenshot API Demo</h1>

    <p>Test the screenshot API with global font support</p>

    

    ${API_KEY ? `

    <div class="auth-section">

        <h3>🔐 API Authentication Required</h3>

        <label>API Key:</label><br>

        <input type="password" id="apiKey" placeholder="Enter your API key" style="width: 400px;">

        <div class="auth-info">

            ℹ️ API key is required for this Space. Contact the Space owner for access.

        </div>

    </div>

    ` : `

    <div class="auth-section">

        <h3>🔓 Open Access Mode</h3>

        <div class="auth-info">

            ℹ️ No API key required - authentication is disabled for this Space.

        </div>

    </div>

    `}

    

    <div class="form-group">

        <label>URL to screenshot:</label><br>

        <input type="text" id="url" value="https://www.baidu.com" placeholder="Enter website URL">

    </div>

    

    <div class="form-group">

        <label>Width:</label>

        <input type="number" id="width" value="1280" min="100" max="1600">

        <label>Height:</label>

        <input type="number" id="height" value="720" min="100" max="1200">

        <label>Quality:</label>

        <input type="number" id="quality" value="75" min="10" max="100">

        <label>Full Page:</label>

        <input type="checkbox" id="fullPage">

        <label>Format:</label>

        <select id="format">

            <option value="jpeg" selected>JPEG</option>

            <option value="png">PNG</option>

        </select>

    </div>

    

    <button onclick="takeScreenshot()">📸 Take Screenshot</button>

    <button onclick="testHealth()">🔍 Health Check</button>

    

    <div id="result"></div>

    

    <script>

        async function testHealth() {

            try {

                const response = await fetch('/');

                const data = await response.json();

                document.getElementById('result').innerHTML = 

                    '<h3>✅ Health Check Result:</h3><pre>' + JSON.stringify(data, null, 2) + '</pre>';

            } catch (error) {

                document.getElementById('result').innerHTML = '<p class="error">❌ Error: ' + error.message + '</p>';

            }

        }

        

        async function takeScreenshot() {

            const url = document.getElementById('url').value;

            const width = parseInt(document.getElementById('width').value);

            const height = parseInt(document.getElementById('height').value);

            const quality = parseInt(document.getElementById('quality').value);

            const fullPage = document.getElementById('fullPage').checked;

            const format = document.getElementById('format').value;

            const apiKey = document.getElementById('apiKey') ? document.getElementById('apiKey').value : '';

            

            if (!url) {

                alert('Please enter a URL');

                return;

            }

            

            ${API_KEY ? `

            if (!apiKey) {

                alert('Please enter your API key');

                return;

            }

            ` : ''}

            

            document.getElementById('result').innerHTML = '<p class="loading">📸 Taking screenshot...</p>';

            

            try {

                const headers = {

                    'Content-Type': 'application/json',

                };

                

                ${API_KEY ? `

                if (apiKey) {

                    headers['Authorization'] = 'Bearer ' + apiKey;

                }

                ` : ''}

                

                const response = await fetch('/screenshot', {

                    method: 'POST',

                    headers: headers,

                    body: JSON.stringify({ url, width, height, quality, fullPage, format })

                });

                

                if (response.ok) {

                    const blob = await response.blob();

                    const imageUrl = URL.createObjectURL(blob);

                    const size = (blob.size / 1024).toFixed(1);

                    

                    document.getElementById('result').innerHTML = 

                        '<h3 class="success">✅ Screenshot Success!</h3>' +

                        '<p>Size: ' + size + ' KB | Dimensions: ' + width + 'x' + height + '</p>' +

                        '<img src="' + imageUrl + '" alt="Screenshot"><br><br>' +

                        '<a href="' + imageUrl + '" download="screenshot.' + format + '">💾 Download Image</a>';

                } else {

                    const error = await response.json();

                    document.getElementById('result').innerHTML = 

                        '<p class="error">❌ Error: ' + (error.error || error.message) + '</p>';

                }

            } catch (error) {

                document.getElementById('result').innerHTML = 

                    '<p class="error">❌ Network Error: ' + error.message + '</p>';

            }

        }

    </script>

</body>

</html>

  `);
});

// 截图端点
app.post('/screenshot', authenticate, async (req, res) => {
  const { url, width = 1280, height = 720, quality = 75, fullPage = false, format = 'jpeg' } = req.body;

  // 验证参数
  if (!url) {
    return res.status(400).json({ error: 'URL is required' });
  }

  if (!url.startsWith('http://') && !url.startsWith('https://')) {
    return res.status(400).json({ error: 'URL must start with http:// or https://' });
  }

  // 验证图片格式
  if (!['jpeg', 'png'].includes(format.toLowerCase())) {
    return res.status(400).json({ error: 'Format must be either "jpeg" or "png"' });
  }

  const cpuUsage = getCpuUsage();
  if (cpuUsage > 95) {
    return res.status(503).json({
      status: 'busy',
      error: 'Server is currently overloaded',
      cpuUsage: `${cpuUsage}%`,
      queueLength: requestQueue.length
    });
  }

  if (activeRequests >= MAX_CONCURRENT) {
    return res.status(503).json({
      status: 'busy',
      error: 'Too many concurrent requests',
      activeRequests,
      maxConcurrent: MAX_CONCURRENT
    });
  }

  activeRequests++;
  let browser;

  try {
    const browserOptions = {
      headless: 'new',
      executablePath: '/usr/bin/google-chrome-stable',
      args: [
        '--no-sandbox',
        '--disable-setuid-sandbox',
        '--disable-dev-shm-usage',
        '--disable-gpu',
        '--no-first-run',
        '--no-zygote',
        '--single-process',
        '--disable-extensions',
        '--disable-background-timer-throttling',
        '--disable-backgrounding-occluded-windows',
        '--disable-renderer-backgrounding',
        // 字体优化设置
        '--font-render-hinting=none',
        '--disable-font-subpixel-positioning',
        '--force-color-profile=srgb',
        '--disable-features=VizDisplayCompositor'
      ]
    };

    browser = await puppeteer.launch(browserOptions);
    const page = await browser.newPage();
    
    // 设置字体回退CSS - 支持全球字体
    await page.addStyleTag({
      content: `

        * {

          font-family: "Noto Sans", "Noto Sans CJK SC", "Noto Sans CJK TC", 

                       "Noto Sans CJK JP", "Noto Sans CJK KR", 

                       "Noto Color Emoji", system-ui, -apple-system, 

                       BlinkMacSystemFont, sans-serif !important;

        }

      `
    });
    
    await page.setViewport({
      width: parseInt(width),
      height: parseInt(height),
      deviceScaleFactor: 1
    });

    await page.goto(url, { 
      waitUntil: 'domcontentloaded', 
      timeout: 15000 
    });

    // 配置截图选项
    const screenshotOptions = {
      type: format.toLowerCase(),
      fullPage: Boolean(fullPage)
    };

    // 只有 JPEG 格式才支持 quality 参数
    if (format.toLowerCase() === 'jpeg') {
      screenshotOptions.quality = parseInt(quality);
    }

    const screenshot = await page.screenshot(screenshotOptions);

    // 设置正确的 Content-Type
    const mimeType = format.toLowerCase() === 'png' ? 'image/png' : 'image/jpeg';
    res.set('Content-Type', mimeType);
    res.send(screenshot);

  } catch (error) {
    console.error('Screenshot error:', error);
    res.status(500).json({ 
      error: 'Failed to capture screenshot', 
      message: error.message 
    });
  } finally {
    if (browser) {
      await browser.close();
    }
    activeRequests--;
  }
});

app.listen(PORT, () => {
  console.log(`🚀 Page Screenshot API running on port ${PORT}`);
  console.log(`📸 Demo: http://localhost:${PORT}/demo`);
  console.log(`🔍 Health: http://localhost:${PORT}/`);
  console.log(`🔐 Authentication: ${API_KEY ? 'ENABLED' : 'DISABLED'}`);
});

module.exports = app;