Windows下pdf文件自动打水印工具

背景：

对指定的pdf文件打上水印，工具包含以下参数功能。

1. 指定输入（pdf文件或路径）。输入既可以是pdf文件，也可以是路径，如果是路径的话，就从路径下面找到所有的pdf文件。

2. 指定输出（路径）。所有的输入pdf文件被打上水印后存放在输出路径下。

3. 指定时间间隔。针对指定的路径，每隔指定的时间就自动去执行一次文件时间戳检测，有任何改动即重新给所有的pdf文件打一次水印。

脚本实现用的python3.6，测试环境为win7，最终实现方案是打包成exe，可以在其他windows机器中随意使用。

下文中的pdf水印文件实现只是个基本方案，我们实际操作是先设计好（公司）logo，然后在打水印时直接用设计好的水印模板打水印。如果想偷懒，如文中直接在pdf底部打印log也可以。

话不多说，直接上代码。

# -*- coding: utf-8 -*-
import os
import re
import sys
import argparse
import time
import collections
from PyPDF2 import PdfFileWriter, PdfFileReader
from reportlab.pdfgen import canvas

# 安装微软雅黑字体（支持中文）。
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
pdfmetrics.registerFont(TTFont('msyh', 'Msyh.ttf'))

cwd=os.getcwd()

def readArgs():
"""
读入参数。
检测参数。
"""
parser = argparse.ArgumentParser()
parser.add_argument('-p', '--pdf',
nargs='+',
default=[],
help='Specify pdf file(s) to add watermark on.',
)
parser.add_argument('-d', '--directory',
default='',
help='Specify directory(s), will find all of the pdf files under them, then add watermark.',
)
parser.add_argument('-o', '--outputDir',
default=str(cwd) + '/watermarkOut',
help='Specify output dir, default is "<CWD>/watermarkOut".',
)
parser.add_argument('-i', '--interval',
type=int,
default=0,
help='Specify the directory monitor interval time (Only if "directory" is specified).')

args = parser.parse_args()
pdfList = []

if (len(args.pdf) == 0) and (args.directory == ''):
# "pdf"和"directory"这两个参数至少指定一个。
print('*Error*: None of "pdf" and "directory" is specified.')
sys.exit(1)
elif (len(args.pdf) > 0) and (args.directory != ''):
# "pdf"和"directory"这两个参数不能同时指定。
print('*Error*: "pdf" and "directory" cannot be specified the same time.')
sys.exit(1)
else:
# 如果"pdf"被指定，且所指定的pdf均存在，将所有pdf文件加入列表pdfList。
if len(args.pdf) > 0:
for pdf in args.pdf:
pdf = pdf.strip()
if os.path.exists(pdf):
pdfList.append(pdf)
else:
print('*Error*: pdf "' + str(pdf) + '": No such file!')
sys.exit(1)

# 如果"directory"被指定，且所指定的directory存在，将找到指定目录下所有的pdf文件加入pdfList。
if args.directory != '':

7fe0
args.directory = args.directory.strip()
if os.path.exists(args.directory):
pdfList = findPdfFiles(args.directory, relativePath='', pdfList=[])
else:
print('*Error*: directory "' + str(args.directory) + '": No such directory!')
sys.exit(1)

# 如果没有任何pdf被指定，报错退出。
if len(pdfList) == 0:
print('*Error*: None pdf file is specified!')
sys.exit(1)

# 如果outputDir不存在，创建它。
if not os.path.exists(args.outputDir):
os.makedirs(args.outputDir)

# "interval"不能小于0。
# "interval"仅当"directory"定义时才起作用。
if args.interval < 0:
print('*Error*: "interval" cannot be less than 0.')
sys.exit(1)

if (args.interval > 0) and (args.directory == ''):
args.interval = 0

return(pdfList, args.directory, args.outputDir, args.interval)

def findPdfFiles(rootPath, relativePath='', pdfList=[]):
"""
在指定的顶层目录下找到所有的pdf文件，输出所有的pdf列表（文件路径为相对于rootDir的相对路径）。
"""
if relativePath == '':
topPath = rootPath
else:
topPath = os.path.join(rootPath, relativePath)

# 遍历顶层目录下所有的项，包括文件和子目录。
items = os.listdir(topPath)

for item in items:
if os.path.isfile(os.path.join(topPath, item)):
# 如果目录下存在pdf文件，直接将pdf文件加入到pdfList。
if re.match('^.*\.pdf$', item):
if relativePath == '':
pdfList.append(item)
else:
pdfList.append(os.path.join(relativePath, item))
elif os.path.isdir(os.path.join(topPath, item)):
# 如果目录下存在子目录，递归调用函数findPdfFiles来获取子目录中的pdf文件。
if relativePath == '':
relativePathSub = item
else:
relativePathSub = os.path.join(relativePath, item)
pdfList = findPdfFiles(rootPath, relativePathSub, pdfList)

return(pdfList)

class watermark():
"""
给指定的pdf文件们添加水印。
"""
def __init__(self, pdfList, inputDir, outputDir, interval):
self.pdfList = pdfList
self.inputDir = inputDir
self.outputDir = outputDir
self.interval = interval

def createWatermark(self, content):
"""
创建PDF水印模板。
"""
# 使用reportlab来创建一个PDF文件来作为一个水印文件。
watermarkTemplate = str(self.outputDir) + '/.watermark.pdf'
c = canvas.Canvas(watermarkTemplate)
c.setFont('msyh', 10)

# 设置水印文件。
c.saveState()
c.translate(300, 15)

# 水印文字。
c.drawCentredString(0, 0, content)
c.restoreState()

# 保存水印文件。
c.save()
pdfWatermark = PdfFileReader(open(watermarkTemplate, 'rb'))

return(pdfWatermark)

def addWatermark(self, pdf, pdfWatermark):
"""
给指定PDF文件文件加上水印。
pdf - 要加水印的源PDF文件。
pdfWatermark - PDF水印模板。
"""
# 获取pdf文件的路径和文件名。
pdfDir = os.path.dirname(pdf)
pdfFile = os.path.basename(pdf)

# 读取pdf文件，获取文件页数。
if self.inputDir == '':
inputPdf = pdf
else:
inputPdf = os.path.join(self.inputDir, pdf)

if not os.path.exists(inputPdf):
print('*Warning*: pdf file "' + str(inputPdf) + '" is missing, cannot add watermark for it.')
return()

inputStream = open(inputPdf, 'rb')
pdfInput = PdfFileReader(inputStream)
pageNum = pdfInput.getNumPages()

pdfOutput = PdfFileWriter()

# 给每一页打水印(将原pdf文件页和水印文件合并到一起)。
for i in range(pageNum):
page = pdfInput.getPage(i)
page.mergePage(pdfWatermark.getPage(0))
pdfOutput.addPage(page)

# 如果输入指定了directory，在outputDir下重建input pdf的目录结构。
if self.inputDir == '':
pdfAbsolutePath = self.outputDir
else:
pdfAbsolutePath = os.path.join(self.outputDir, pdfDir)
if not os.path.exists(pdfAbsolutePath):
os.makedirs(pdfAbsolutePath)

# 输出文件。
outputPdf = os.path.join(pdfAbsolutePath, pdfFile)
if os.path.exists(outputPdf):
try:
os.remove(outputPdf)
except Exception as warning:
print('*Warning*: Failed on removing file "' + str(outputPdf) + '": ' + str(warning))

try:
outputStream = open(outputPdf, 'wb')
pdfOutput.write(outputStream)
outputStream.close()
inputStream.close()
except Exception as warning:
print('*Warning*: Failed on open "' + str(outputPdf) + '" for write: ' + str(warning))
outputStream.close()
inputStream.close()

def getTimeStamp(self):
"""
给出一个文件列表，找到每个文件的修改时间戳，并把(文件：修改时间戳)保存在字典中并返回。
"""
timeStampDic = collections.OrderedDict()

for pdf in self.pdfList:
if self.inputDir == '':
pdfFile = pdf
else:
pdfFile = os.path.join(self.inputDir, pdf)

if os.path.exists(pdfFile):
timeStemp = os.stat(pdfFile).st_mtime
timeStampDic[pdf] = timeStemp
else:
timeStampDic[pdf] = 0

return(timeStampDic)

def compareTimeStamp(self, timeStampDic1, timeStampDic2):
"""
给定两个时间戳字典，若两个字典的元素不同，或元素相同但某个文件时间戳不同，返回False，否则返回True。
"""
fileList1 = list(timeStampDic1.keys())
fileList2 = list(timeStampDic2.keys())

if fileList1 == fileList2:
for file in fileList1:
timeStamp1 = timeStampDic1[file]
timeStamp2 = timeStampDic2[file]
if timeStamp1 != timeStamp2:
return(False)
else:
return(False)

return(True)

def watermark(self):
"""
给指定的pdf文件们打上水印，输出到指定路径。
"""
pdf_watermark = self.createWatermark('LYQ 版权所有')
timeStampDic = self.getTimeStamp()
timeStampDicOld = collections.OrderedDict()

while True:
if not self.compareTimeStamp(timeStampDic, timeStampDicOld):
timeStampDicOld = timeStampDic
for pdf in self.pdfList:
print('Add watermark for "' + str(pdf) + '"')
self.addWatermark(pdf, pdf_watermark)

if (self.inputDir != '') and (self.interval > 0):
print('Sleeping ' + str(self.interval) + ' seconds ...')
time.sleep(self.interval)
self.pdfList = findPdfFiles(self.inputDir, relativePath='', pdfList=[])
timeStampDic = self.getTimeStamp()
else:
break

#################
# Main Function #
#################
def main():
(pdfList, inputDir, outputDir, interval) = readArgs()
myWaterMark = watermark(pdfList, inputDir, outputDir, interval)
myWaterMark.watermark()

if __name__ == '__main__':
main()