Python 为PDF 文件添加水印和加密

1,012 阅读3分钟

携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第25天,点击查看活动详情

前言

为了更好的保证数据安全,在 PDF 文件中可以使用一些安全措施来限制文档的泄露。我们可以对内容进行加密,要求用户输入密码才能阅读,或者我们也可以通过添加水印以清楚地将文档标记为不用于公开传播,并且在发生泄露时,了解泄露发生的原因。

Python 为PDF 文件添加水印和加密

在本节中,我们将使用 pdf2image 模块将 PDF 文档转换为 PIL 图像,为了完成这一目标我们还需要首先安装 pillow 库,除此之外,我们还需要使用 PyPDF2

为了完成加密和添加水印任务,我们首先编写 Python 脚本 encrypt_watermarking.py

import argparse
import os
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont
import PyPDF2
from pdf2image import convert_from_path

DEFAULT_OUTPUT = 'watermarked.pdf'
DEFAULT_BY = 'default user'
INTERMEDIATE_ENCRYPT_FILE = 'temp.pdf'

WATERMARK_SIZE = (200, 200)

def encrypt(out_pdf, password):
    print('Encrypting the document')

    output_pdf = PyPDF2.PdfFileWriter()

    in_file = open(out_pdf, "rb")
    input_pdf = PyPDF2.PdfFileReader(in_file)
    output_pdf.appendPagesFromReader(input_pdf)
    output_pdf.encrypt(password)

    # Intermediate file
    with open(INTERMEDIATE_ENCRYPT_FILE, "wb") as out_file:
        output_pdf.write(out_file)

    in_file.close()

    # Rename the intermediate file
    os.rename(INTERMEDIATE_ENCRYPT_FILE, out_pdf)


def create_watermark(watermarked_by):
    print('Creating a watermark')
    mask = Image.new('L', WATERMARK_SIZE, 0)
    draw = ImageDraw.Draw(mask)
    font = ImageFont.load_default()
    text = 'WATERMARKED BY {}\n{}'.format(watermarked_by, datetime.now())
    draw.multiline_text((0, 100), text, 55, font=font)

    watermark = Image.new('RGB', WATERMARK_SIZE)
    watermark.putalpha(mask)
    watermark = watermark.resize((1950, 1950))
    watermark = watermark.rotate(45)
    # Crop to only the watermark
    bbox = watermark.getbbox()
    watermark = watermark.crop(bbox)

    return watermark


def apply_watermark(watermark, in_pdf, out_pdf):
    print('Watermarking the document')
    # Transform from PDF to images
    images = convert_from_path(in_pdf)

    # Get the location for the watermark
    hi, wi = images[0].size
    hw, ww = watermark.size
    position = ((hi - hw) // 2, (wi - ww) // 2)

    # Paste the watermark in each page
    for image in images:
        image.paste(watermark, position, watermark)

    # Save the resulting PDF
    images[0].save(out_pdf, save_all=True, append_images=images[1:])


def main(in_pdf, out_pdf, watermarked_by, password):

    watermark = create_watermark(watermarked_by)
    apply_watermark(watermark, in_pdf, out_pdf)

    if password:
        encrypt(out_pdf, password)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(dest='pdf', type=str, help='PDF to watermark')
    parser.add_argument('-o', type=str, help=f'Output PDF filename, default: {DEFAULT_OUTPUT}', default=DEFAULT_OUTPUT)
    parser.add_argument('-u', type=str, help=f'Watermarked by, default: {DEFAULT_BY}', default=DEFAULT_BY)
    parser.add_argument('-p', type=str, help='Password')
    args = parser.parse_args()
    main(args.pdf, args.o, args.u, args.p)

encrypt_watermarking.py 脚本首先使用 argparse 从命令行获取参数,然后将其传递给 main 函数,该函数负责调用函数 create_watermarkapply_watermark,同时如果需要使用密码,则会调用 password 函数进行加密。

create_watermark 函数用于生成带有水印的图像。它首先使用 Pillow Image 类创建灰度图像(模式 L )并绘制文本。然后,此图像将作为 Alpha 通道应用于 PDF 文档图像,使用 Alpha 通道可以使图像半透明,以便将文本显示为水印。

Alpha 通道使白色(颜色 0 )值的像素区域完全透明,黑色(颜色 255 )值的像素区域完全不透明。为了不影响原始 PDF 文档的阅读性,我们设定背景为白色(完全透明),文本的颜色为 55,使其成为半透明。

然后将图像旋转 45 度并裁剪以适应 PDF 文档,以使图像居中并进行更好的定位。

apply_watermark 函数使用 pdf2image 模块将 PDF 转换为 PIL 图像序列,然后计算应用水印的位置,然后将在 create_watermark 函数中生成水印图像粘贴在 PDF 图像序列上。

如果在命令行中使用 -p 选项传递了密码,则调用 encrypt 函数。它使用 PdfFileReader 打开输出 PDF,并使用 PdfFileWriter 创建一个新的临时 PDF,然后,输出 PDF 的所有页面添加到新 PDF 中,对 PDF 进行加密,然后使用 os.rename 将中间 PDF 重命名为新的 PDF 文档。

编写完成后,可以使用以下命令为 PDF 文件添加水印:

$ python encrypt_watermarking.py document.pdf -u User_panxiaohui -o output.pdf

其中 encrypt_watermarking.py 是编写完成的程序脚本,document.pdf 是待添加水印的 PDF 文档,-u 用于指定添加水印的内容,此处使用用户名 User_panxiaohui-o 用于指定添加水印后保存的文件名,此处使用 output.pdf

检查文档是否在 output.pdf 的所有页面上添加了带有 User_panxiaohui 和时间戳的水印:

添加水印结果

可以使用以下命令同时为 PDF 文档添加水印和加密,加密可能需要等待一定的时间:

$ python encrypt_watermarking.py document.pdf -u User_panxiaohui -o new_output.pdf -p password000

在以上命令中,大部分参数与添加水印使用的命令完全相同,增加的 -p 选项用于加密 PDF 文档,本例中使用值 password000 作为密码。

打开生成的 new_output.pdf 文件并检查它是否需要输入 password000 密码才能访问。