Python | 批量生成加密PDF文件

做数据的二号姬

发布于 2025-01-07 14:52:35

8800

代码可运行

文章被收录于专栏：做数据的二号姬做数据的二号姬

运行总次数：0

代码可运行

年底年初照例是比较忙的时候，最近捣鼓了一个小的脚本：读取数据库员工绩效数据，批量生成加密码加水印的业绩数据PDF文件。

虽然我个人是第一次遇到这个续期，但总的来说对于数分而言，这还是泛用性比较高的需求场景，所以稍微多研究了一下。

捣鼓这个脚本前后零零散散花了比较多的时间，主要是在折腾AI辅助编程相关的事情，对于调教AI写代码这件事情上也算是比较有心得。我准备组织多篇文章给大家分享其中有意思的细节，今天就先介绍一下脚本本身。

前半部分我会给到代码的思路和讲解，完整的代码放在最后，没有时间阅读细节的朋友们可以直接拉到最后看完整的代码。

绝大多数数据分析师多少都能写点代码，但毕竟我不是专业的程序员，写这种小脚本其实还是有点困难的。毕竟数据分析师很多都是主要使用jupyter notebook做数据处理和计算，这种批量生成加密加水印的PDF文件这种小需求对于程序员来说可能没啥，但是对于数据分析师来说多少还是有一点难度的。

我第一次做类似小需求的时候也是一脸迷茫，对于编程的思路完全不理解，做多了之后就稍微有点心得，虽然都是写Python代码，这种小需求和pandas处理数据从思路上是完全不一样的，虽然简单但确确实实是应该应用编程的逻辑，所以我就简单分享一下思路方面的事情。

对于这种小需求，我们首先要对需求本身进行拆解，基本可以分为几个部分：①读取数据库获取员工绩效数据；②根据员工的工号将数据分组并保存为PDF文件；③给PDF文件添加水印和密码。

因为最终的目标是批量生成，所以我们是思路就是做遍历，做法步骤上有很多种解题思路：

第一种做法是：读取第一个人的数据生成加密加水印的PDF，然后读取第二个人的数据生成加密加水印的PDF文件，然后读取第三个人的……

第二种做法是：我们先把所有人的数据都读取出来，然后先生成第一个的加密加水印PDF文件，再生成第二个人的，再生成第三个人的……

第三种做法是：我们先把所有人的数据都读取过来，批量创建PDF文件，再批量打水印加密码……

这个问题比较简单，实际上我们在构思程序框架的时候也并不需要把所有的实现方式都罗列出来，因为至少稍微写出一两种方式就很容易能找出其中的共性：执行SQL获取数据、生成PDF文件、PDF文件加密码、PDF文件加水印。

做数分的同学应该比较好理解，编程的一大要诀就是把这些共性的东西写成函数。如果选择第二种方法的话，上面我们已经把框架里出来了，所以代码最终的框架就有了：

思路已经有了，接下来就是coding了。

第一步就是要安装需要的第三方库，并且导入包，主要需要用到的三方库有两个，reportlib库和pypdf2。没有安装的可以考虑用下面的方式进行安装：

pip install reportlib
pip install pypdf2

reportlib库主要用来生成PDF文件，pypdf2主要用来对PDF文件进行加密处理。此外，因为数据库是mysql外加我需要把加密PDF的密码也另存一份，所以还导入了mysql和openpyxl。

import mysql.connector
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.lib.colors import Color
from reportlab.lib import colors
from PyPDF2 import PdfWriter, PdfReader
import openpyxl
import os
from datetime import datetime
import random
import string
from reportlab.platypus import BaseDocTemplate, PageTemplate, Frame, SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle

接下来就是连接数据库访问数据，实际上这部分内容比较简单，不同类型的数据库差异比较大，因为我选择的是上面的思路二，而且大概率是个一次性需求，所以其实没有必要def一个函数，具体需要不需要还是要看大家实际的需求。


# 数据库连接配置
db_config = {
    'user': '用户名',
    'password': '密码',
    'host': '数据库地址',
    'database': '数据库',
    'raise_on_warnings': True,
    'port': '端口号'
}

# 连接数据库
cnx = mysql.connector.connect(**db_config)
cursor = cnx.cursor()

# 执行 SQL 查询
query = '''
SQL语句
'''

cursor.execute(query)

# 获取列名
column_names = [desc[0] for desc in cursor.description]

rows = cursor.fetchall()
# 将元组转换为字典列表
dcrs = [dict(zip(column_names, row)) for row in rows]

# 关闭数据库连接
cursor.close()
cnx.close()

因为最后要by员工批量生成数据，所以这里需要对数据进行分组操作。

grouped_data = {}
for row in dcrs:
    employee_cd = row["分组的字段名"]
    if employee_cd not in grouped_data:
        grouped_data[employee_cd] = []
    grouped_data[employee_cd].append(row)

接下来就是创建PDF文件了，首先，我的文档内容都是中文的，我得先注册一个中文字体，防止数据乱码（或许不能说防止乱码，因为不注册100%乱码）。

# 注册中文字体，这里以微软雅黑为例
pdfmetrics.registerFont(TTFont('Microsoft YaHei', 'msyh.ttc'))
Style = getSampleStyleSheet()
pStyle = Style['Normal']

我在直接创建PDF的过程中遇到了一点小问题，那就是水印和页码只对第一页生效，所以单独创建了一个class用于实现水印和页码在每页都生效。

class MyDocTemplate(BaseDocTemplate):
    def __init__(self, filename, watermark_text, **kw):
        self.watermark_text = watermark_text
        BaseDocTemplate.__init__(self, filename, **kw)
        frame = Frame(self.leftMargin, self.bottomMargin, self.width, self.height, id='normal')
        self.addPageTemplates(
            [PageTemplate(id='test', frames=frame, onPage=self.add_watermark, onPageEnd=self.add_page_footer)])

    #创建水印
    def add_watermark(self, canvas, doc):
        canvas.saveState()
        canvas.setFillColorRGB(0.5, 0.5, 0.5, alpha=0.3)
        canvas.setFont('Microsoft YaHei', 18)
        angle = 45
        width, height = letter
        delta_x = width / 3
        delta_y = height / 5

        # 修改水印位置
        for x in range(0, round(width), round(delta_x)):
            for y in range(0, round(height), round(delta_y)):  
                canvas.saveState()
                canvas.translate(x, y)
                canvas.rotate(angle)
                canvas.drawString(-100, -150, self.watermark_text)
                canvas.restoreState()
        canvas.restoreState()

    #添加页脚，页码
    def add_page_footer(self, canvas, doc):
        canvas.saveState()
        style = ParagraphStyle(name='SmallText', fontSize=8, fontName='Microsoft YaHei')
        page = doc.page
        total_pages = doc.page
        footer_text = Paragraph(f'第{page}页', style)
        footer_text.wrapOn(canvas, letter[0], letter[1])
        footer_text.drawOn(canvas, 300, 28)
        canvas.restoreState()

创建水印主要是用了canvas方法，水印的内容是员工的工号，所以写成了一个外部传入的参数。这里稍微有一个小问题，我原本想设置第X页，共X页的，但是共X页有点问题，实现不了，就放弃了，不是很大的影响。

接下来就是创建pdf文件了。为了实现表格自动分页且重复表头，这里使用了MyDocTemplate进行创建，相对来说比canvas简单一些。这种方法就是一个模板，结合上面的class，实际上就是先创建了一个文件的模板，然后把内容写入。

def create_pdf(data, filename, watermark_text):
    #指定支持中文
    styles = getSampleStyleSheet()
    styleN = styles['Normal']
    styleH1 = styles['Heading1']
    styleH2 = styles['Heading2']
    styleH1.fontName = 'Microsoft YaHei'
    styleH2.fontName = 'Microsoft YaHei'

    #创建文件

    doc = MyDocTemplate(filename, watermark_text, pagesize=letter, leftMargin=36,
                        rightMargin=36, topMargin=36, bottomMargin=36)

    elements = []

    elements.append(Paragraph("一级标题", styleH1))
    elements.append(Spacer(1, 12))

    elements.append(Paragraph(
        f'Employee_Code: {data[0].get("employee_cd", "N/A")} 报告日期: {datetime.now().strftime("%Y年%m月%d日")}', styleH2))
    elements.append(Spacer(1, 12))


    headers = ['数据表头1'，'数据表头2']
    pStyle = ParagraphStyle(name='ChineseStyle', fontSize=10, fontName='Microsoft YaHei')

    table_data = [headers]
    for row in data:
        row_data = [
            Paragraph(row["字段名1"], pStyle),
            Paragraph(row["字段名2"], pStyle),
          
        ]
        table_data.append(row_data)

    #创建表格，跨页时重复表头
    table = Table(table_data, repeatRows=1)

    #设置表格样式细节
    table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
        ('FONTNAME', (0, 0), (-1, 0), 'Microsoft YaHei'),
        ('FONTNAME', (0, 1), (-1, -1), 'Microsoft YaHei'),
        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
        ('BACKGROUND', (0, 1), (-1, -1), colors.transparent),  #设置表格背景色为透明
        ('BOX', (0, 0), (-1, -1), 0.25, colors.black),
        ('GRID', (0, 0), (-1, -1), 0.5, colors.black),
        ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
        ('ROWHEIGHT', (0, 0), (-1, -1), 40)
    ]))

    elements.append(table)

    doc.build(elements)

这里稍微有一个小问题，那就是水印写在文档模板里，其实是在表格下方的，但是表格的背景色我设置成了透明，就不影响了。根据我的研究，如果想把水印搞到上面来就得用canvas方法生成表格，这种方法对于跨页自动换行没有支持，自己写很麻烦了，我尝试了很久页没成功。

接下来就相对简单了，def一个函数对PDF文件进行加密：

# 加密PDF文件
def encrypt_pdf(input_pdf, output_pdf, password):
    pdf_writer = PdfWriter()
    pdf_reader = PdfReader(input_pdf)

    for page in range(len(pdf_reader.pages)):
        pdf_writer.add_page(pdf_reader.pages[page])
        pdf_writer.encrypt(user_password=password, owner_pwd=None, use_128bit=True)

    with open(output_pdf, 'wb') as out:
        pdf_writer.write(out)

因为文件要加密，所以还需要生成随机密码。

# 生成随机密码
def generate_random_password(length=16):
    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

把生成的随机密码写入excel表里：

# 将密码写入Excel
def write_passwords_to_excel(filenames_passwords, excel_filename):
    workbook = openpyxl.Workbook()
    sheet = workbook.active
    sheet['A1'] = 'File Name'
    sheet['B1'] = 'Password'
    for i, (filename, password) in enumerate(filenames_passwords, start=2):
        sheet[f'A{i}'] = filename
        sheet[f'B{i}'] = password
        workbook.save(excel_filename)

接下来就是执行代码：

# 遍历grouped_data，生成对应的pdf文件，并加密
def generate_pdfs_and_encrypt(grouped_data):
    filenames_passwords = []
    for group_key, data in grouped_data.items():
        watermark_text = data[0]['分组字段'] if data else 'No Data'
        pdf_filename = f'{group_key}.pdf'
        create_pdf(data, pdf_filename, watermark_text)
        password = generate_random_password()
        encrypted_pdf_filename = f'{group_key}.pdf'
        encrypt_pdf(pdf_filename, encrypted_pdf_filename, password)
        filenames_passwords.append((encrypted_pdf_filename, password))
        os.remove(pdf_filename)  # 删除未加密的PDF文件
    return filenames_passwords


filenames_passwords = generate_pdfs_and_encrypt(grouped_data)

# 将文件名和密码写入Excel
excel_filename = 'pdf_passwords.xlsx'
write_passwords_to_excel(filenames_passwords, excel_filename)

print(f"所有PDF文件已加密，密码已写入Excel文件: {excel_filename}")

然后就没有然后了，虽然代码本身不算很复杂，但是我折腾了很久，很大程度上是因为AI误导给我带坑里了，关于调教AI的趣事，下篇文章继续说。

最后是完整版本的代码，供参考

import mysql.connector
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.lib.colors import Color
from reportlab.lib import colors
from PyPDF2 import PdfWriter, PdfReader
import openpyxl
import os
from datetime import datetime
import random
import string
from reportlab.platypus import BaseDocTemplate, PageTemplate, Frame, SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
# 数据库连接配置
db_config = {
    'user': '用户名',
    'password': '密码',
    'host': '数据库地址',
    'database': '数据库',
    'raise_on_warnings': True,
    'port': '端口号'
}
# 连接数据库
cnx = mysql.connector.connect(**db_config)
cursor = cnx.cursor()
# 执行 SQL 查询
query = '''
SQL语句
'''
cursor.execute(query)
# 获取列名
column_names = [desc[0] for desc in cursor.description]
rows = cursor.fetchall()
# 将元组转换为字典列表
dcrs = [dict(zip(column_names, row)) for row in rows]
# 关闭数据库连接
cursor.close()
cnx.close()
# 根据 员工号进行 分组
grouped_data = {}
for row in dcrs:
    employee_cd = row["分组的字段名"]
    if employee_cd not in grouped_data:
        grouped_data[employee_cd] = []
    grouped_data[employee_cd].append(row)
# 注册中文字体，这里以微软雅黑为例
pdfmetrics.registerFont(TTFont('Microsoft YaHei', 'msyh.ttc'))
Style = getSampleStyleSheet()
pStyle = Style['Normal']
class MyDocTemplate(BaseDocTemplate):
    def __init__(self, filename, watermark_text, **kw):
        self.watermark_text = watermark_text
        BaseDocTemplate.__init__(self, filename, **kw)
        frame = Frame(self.leftMargin, self.bottomMargin, self.width, self.height, id='normal')
        self.addPageTemplates(
            [PageTemplate(id='test', frames=frame, onPage=self.add_watermark, onPageEnd=self.add_page_footer)])
    #创建水印
    def add_watermark(self, canvas, doc):
        canvas.saveState()
        canvas.setFillColorRGB(0.5, 0.5, 0.5, alpha=0.3)
        canvas.setFont('Microsoft YaHei', 18)
        angle = 45
        width, height = letter
        delta_x = width / 3
        delta_y = height / 5
        # 修改水印位置
        for x in range(0, round(width), round(delta_x)):
            for y in range(0, round(height), round(delta_y)):  
                canvas.saveState()
                canvas.translate(x, y)
                canvas.rotate(angle)
                canvas.drawString(-100, -150, self.watermark_text)
                canvas.restoreState()
        canvas.restoreState()
    #添加页脚，页码
    def add_page_footer(self, canvas, doc):
        canvas.saveState()
        style = ParagraphStyle(name='SmallText', fontSize=8, fontName='Microsoft YaHei')
        page = doc.page
        total_pages = doc.page
        footer_text = Paragraph(f'第{page}页', style)
        footer_text.wrapOn(canvas, letter[0], letter[1])
        footer_text.drawOn(canvas, 300, 28)
        canvas.restoreState()
def create_pdf(data, filename, watermark_text):
    #指定支持中文
    styles = getSampleStyleSheet()
    styleN = styles['Normal']
    styleH1 = styles['Heading1']
    styleH2 = styles['Heading2']
    styleH1.fontName = 'Microsoft YaHei'
    styleH2.fontName = 'Microsoft YaHei'
    #创建文件
    doc = MyDocTemplate(filename, watermark_text, pagesize=letter, leftMargin=36,
                        rightMargin=36, topMargin=36, bottomMargin=36)
    elements = []
    elements.append(Paragraph("一级标题", styleH1))
    elements.append(Spacer(1, 12))
    elements.append(Paragraph(
        f'Employee_Code: {data[0].get("employee_cd", "N/A")} 报告日期: {datetime.now().strftime("%Y年%m月%d日")}', styleH2))
    elements.append(Spacer(1, 12))
    headers = ['数据表头1'，'数据表头2']
    pStyle = ParagraphStyle(name='ChineseStyle', fontSize=10, fontName='Microsoft YaHei')
    table_data = [headers]
    for row in data:
        row_data = [
            Paragraph(row["字段名1"], pStyle),
            Paragraph(row["字段名2"], pStyle),
          
        ]
        table_data.append(row_data)
    #创建表格，跨页时重复表头
    table = Table(table_data, repeatRows=1)
    #设置表格样式细节
    table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (-1, 0), colors.grey),
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
        ('FONTNAME', (0, 0), (-1, 0), 'Microsoft YaHei'),
        ('FONTNAME', (0, 1), (-1, -1), 'Microsoft YaHei'),
        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
        ('BACKGROUND', (0, 1), (-1, -1), colors.transparent),  #设置表格背景色为透明
        ('BOX', (0, 0), (-1, -1), 0.25, colors.black),
        ('GRID', (0, 0), (-1, -1), 0.5, colors.black),
        ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
        ('ROWHEIGHT', (0, 0), (-1, -1), 40)
    ]))
    elements.append(table)
    doc.build(elements)
# 加密PDF文件
def encrypt_pdf(input_pdf, output_pdf, password):
    pdf_writer = PdfWriter()
    pdf_reader = PdfReader(input_pdf)
    for page in range(len(pdf_reader.pages)):
        pdf_writer.add_page(pdf_reader.pages[page])
        pdf_writer.encrypt(user_password=password, owner_pwd=None, use_128bit=True)
    with open(output_pdf, 'wb') as out:
        pdf_writer.write(out)
# 将密码写入Excel
def write_passwords_to_excel(filenames_passwords, excel_filename):
    workbook = openpyxl.Workbook()
    sheet = workbook.active
    sheet['A1'] = 'File Name'
    sheet['B1'] = 'Password'
    for i, (filename, password) in enumerate(filenames_passwords, start=2):
        sheet[f'A{i}'] = filename
        sheet[f'B{i}'] = password
        workbook.save(excel_filename)
# 生成随机密码
def generate_random_password(length=16):
    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))
# 遍历grouped_data，生成对应的pdf文件，并加密
def generate_pdfs_and_encrypt(grouped_data):
    filenames_passwords = []
    for group_key, data in grouped_data.items():
        watermark_text = data[0]['分组字段'] if data else 'No Data'
        pdf_filename = f'{group_key}.pdf'
        create_pdf(data, pdf_filename, watermark_text)
        password = generate_random_password()
        encrypted_pdf_filename = f'{group_key}.pdf'
        encrypt_pdf(pdf_filename, encrypted_pdf_filename, password)
        filenames_passwords.append((encrypted_pdf_filename, password))
        os.remove(pdf_filename)  # 删除未加密的PDF文件
    return filenames_passwords
filenames_passwords = generate_pdfs_and_encrypt(grouped_data)
# 将文件名和密码写入Excel
excel_filename = 'pdf_passwords.xlsx'
write_passwords_to_excel(filenames_passwords, excel_filename)

print(f"所有PDF文件已加密，密码已写入Excel文件: {excel_filename}")

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2025-01-05，如有侵权请联系 cloudcommunity@tencent.com 删除

数据