ColdVerify_server/script/add_signature.py

import io
import logging
import os
import sys

import pdfplumber
import requests
from PIL import Image
from PyPDF2 import PdfReader, PdfWriter
from reportlab.lib.pagesizes import A4
from reportlab.pdfgen import canvas

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def find_signature_positions(pdf_path, target_texts=["签字或盖章", "技术支持单位（盖章）"]):
    """查找所有签名位置"""
    positions = []
    with pdfplumber.open(pdf_path) as pdf:
        for i, page in enumerate(pdf.pages):
            for word in page.extract_words():
                for target in target_texts:
                    if target in word["text"]:
                        x, y = word["x0"], word["top"]
                        positions.append((i, int(x), int(y), target))
                        logger.info(f"找到签名位置: 第{i}页, {target} at ({x}, {y})")

    # 返回结果：所有"签字或盖章"和"技术支持单位（盖章）"的位置
    if len(positions) == 0 or positions == []:
        return None
        # 返回结果：第二个"签字或盖章"和技术支持单位位置
    result = []
    if len(positions) >= 2:
        result.append(positions[1])
        # 第二个出现的位置（索引1）
    else:
        result.append(positions[0])

    return result


def add_signature_to_pdf(input_pdf, output_pdf, signature_img, positions, offset_x=30, offset_y=-10):
    """添加签名到PDF"""
    reader = PdfReader(input_pdf)
    writer = PdfWriter()
    a4_width, a4_height = A4  # (595.2, 841.68)
    for page_index in range(len(reader.pages)):
        page = reader.pages[page_index]
        for pos in positions:
            if pos[0] == page_index:
                packet = io.BytesIO()
                can = canvas.Canvas(packet, pagesize=A4)
                img = Image.open(signature_img)
                img_width, img_height = img.size
                aspect_ratio = img_height / img_width
                new_width = 43 * (72 / 25.4)  # 40mm ≈ 113.39点
                new_height = new_width * aspect_ratio
                x, y = pos[1], a4_height - pos[2]
                if pos[3] == "技术支持单位（盖章）":
                    y = y - new_height / 2
                else:
                    y = y - new_height / 3 * 2
                x = max(0, min(x, a4_width - new_width))
                y = max(0, min(y, a4_height - new_height))
                can.drawImage(signature_img, x, y, width=new_width, height=new_height, mask='auto')
                can.save()
                packet.seek(0)
                overlay_pdf = PdfReader(packet)
                page.merge_page(overlay_pdf.pages[0])
        writer.add_page(page)

    with open(output_pdf, "wb") as f:
        writer.write(f)


def download_file(url, local_filename):
    """下载文件到临时目录"""
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(local_filename, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    return local_filename


def add_signature(url, pdfFilename, signature_img, signature_pdf):
    """把水印添加到pdf中"""

    # 下载PDF文件
    local_pdf = download_file(url, pdfFilename)
    # 查找签名位置（使用默认公章图片）
    positions = find_signature_positions(local_pdf)
    if len(positions) == 0:
        os.remove(local_pdf)
        return

    # 生成带公章的PDF
    add_signature_to_pdf(local_pdf, signature_pdf, signature_img, positions)


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("Usage: add_watermark.py <url> <pdf_file_in> <signature_img> <signature_pdf>")
        sys.exit(1)

    pdf_file_url = sys.argv[1]
    pdfFilename = sys.argv[2]  # 要加水印的文件名
    signature_img = sys.argv[3]  # 要加水印的文件名
    signature_pdf = sys.argv[4]  # 加好水印的结果文件

    # name = str(uuid.uuid4()) + ".pdf"
    # pdf_file_url = "https://bzdcoldverifyoss.baozhida.cn/UpImage/1754036272edee05b2-6032-430c-a39b-c83490764dd9.pdf"
    # pdfFilename = "/Users/zoie/work/bzd_project/ColdVerify_server/ofile/" + name
    # signature_img = "/Users/zoie/work/bzd_project/ColdVerify_server/script/报告专用章.png"
    # signature_pdf = "/Users/zoie/work/bzd_project/ColdVerify_server/ofile/signature" + name

    add_signature(pdf_file_url, pdfFilename, signature_img, signature_pdf)
    print(signature_pdf)