import os
import reportlab as reportlab
from PyPDF2 import PdfFileReader, PdfFileWriter
from reportlab.pdfgen import canvas  # 提供pdf文件
from reportlab.pdfbase import pdfmetrics  # 注册字体
from reportlab.pdfbase.ttfonts import TTFont  # 提供字体对象


def water():
    pdfmetrics.registerFont(TTFont('字体1', '青呱石头体.ttf'))
    pdf = canvas.Canvas('water.pdf')
    pdf.setFont('字体1', 60)
    pdf.setFillColorRGB(0.3, 0.4, 0.5, 0.5)
    pdf.rotate(45)
    pdf.drawString(350, 0, '张小海')
    pdf.save()


def add_watermark(file, water):
    water_reader = PdfFileReader(open(water, 'rb'))
    file_reader = PdfFileReader(open(f'pdfs/{file}', 'rb'))
    water_page = water_reader.getPage(0)
    writer = PdfFileWriter()
    a = file[0:-4]
    for i in range(file_reader.numPages):
        page = file_reader.getPage(i)
        page.mergePage(water_page)
        writer.addPage(page)
    writer.write(open(f'{a}水印.pdf', 'wb'))


water()
for i in os.listdir('pdfs'):
    if i.endswith('.pdf'):
        add_watermark(i, 'water.pdf')