给文件夹中的pdf加水印升级版(支持zip,tar.gz)

import os
from unrar import rarfile
from PyPDF2 import PdfFileReader, PdfFileWriter
from reportlab.lib.units import cm
from reportlab.pdfgen import canvas
import zipfile
import tarfile
import shutil
import random


class watermark_sv(object):

    def __init__(self, path, user_name):
        self.six_code = self.six_random()
        # 获取上一级目录
        pre_path = os.path.abspath(os.path.join(os.getcwd(), ".."))
        new_path = os.path.join(pre_path, 'repository/scripts/' + self.six_code)
        if not os.path.exists(new_path):
            os.mkdir(new_path)
        shutil.copy(path, new_path)
        (file_path, file_name) = os.path.split(path)
        print(file_path, file_name)
        # 解压包所在的新路径
        new_path = os.path.join(new_path, file_name)
        # 解压文件
        self.un_zip(new_path)
        # 遍历文件夹给PDF加水印
        self.dir_scripts(os.getcwd())
        #  删除无关的文件
        self.delete_file(os.getcwd())
        # 在scripts目录下进行压缩
        if '.zip' in file_name:
            self.zip_file_path(file_name)
        elif '.gz' in file_name:
            self.tar_gz_file_path(file_name)
        all_files = os.listdir(os.getcwd())
        for file in all_files:
            path = os.path.join(os.getcwd(), file)
            if os.path.isdir(path):
                shutil.rmtree(path)

    # 水印信息
    def create_watermark(self, content):
        # 默认大小为21cm*29.7cm
        file_name = "mark.pdf"
        c = canvas.Canvas(file_name, pagesize=(30 * cm, 30 * cm))
        # 移动坐标原点(坐标系左下为(0,0))
        c.translate(1 * cm, 0 * cm)
        # 设置字体
        c.setFont("Helvetica", 20)
        # 指定描边的颜色
        c.setStrokeColorRGB(0, 1, 0)
        # 指定填充颜色
        c.setFillColorRGB(0, 1, 0)
        # 旋转45度,坐标系被旋转
        c.rotate(45)
        # 指定填充颜色
        c.setFillColorRGB(0, 0, 0, 0.1)
        # 画几个文本,注意坐标系旋转的影响
        for x in range(0, 45, 5):
            for y in range(-20, 22, 2):
                c.drawString(x * cm, y * cm, content)
        # 关闭并保存pdf文件
        c.save()
        return file_name


    # 把水印添加到pdf中
    def add_watermark(self, pdf_file_in, pdf_file_mark, pdf_file_out):
        pdf_output = PdfFileWriter()
        input_stream = open(pdf_file_in, 'rb')
        pdf_input = PdfFileReader(input_stream, strict=False)

        # 获取PDF文件的页数
        pageNum = pdf_input.getNumPages()

        # 读入水印pdf文件
        pdf_watermark = PdfFileReader(open(pdf_file_mark, 'rb'), strict=False)
        # 给每一页打水印
        for i in range(pageNum):
            page = pdf_input.getPage(i)
            page.mergePage(pdf_watermark.getPage(0))
            page.compressContentStreams()  # 压缩内容
            pdf_output.addPage(page)

        # 删除原有的pdf
        path = os.path.join(os.getcwd(), pdf_file_in)
        if os.path.exists(path):
            os.remove(path)
        pdf_output.write(open(pdf_file_out, 'wb'))


    # 删除压缩文件和水印文件
    def delete_file(self, path):
        mark_path = os.path.join(path, 'mark.pdf')
        if os.path.exists(mark_path):
            os.remove(mark_path)
        for list in os.listdir(path):
            path_zip = os.path.join(path, list)
            if '.zip' in path_zip:
                os.remove(path_zip)
            if '.gz' in path_zip:
                os.remove(path_zip)
            if os.path.isdir(path_zip):
                self.delete_file(path_zip)

    # 层级解压zip文件
    def unzip(self, path, zfile):
        file_path = path + os.sep + zfile
        desdir = path + os.sep + zfile[:zfile.index('.zip')]
        srcfile = zipfile.ZipFile(file_path)
        for filename in srcfile.namelist():
            srcfile.extract(filename, desdir)
            if filename.endswith('.zip'):
                path = desdir
                zfile = filename
                self.unzip(path, zfile)
        srcfile.close()

    # 层级解压tar.gz文件
    def untar_dir(self, dstPath, srcname):
        file_path = dstPath + os.sep + srcname
        desdir = dstPath + os.sep + srcname[:srcname.index('.tar.gz')]
        tarHandle = tarfile.open(file_path, "r:gz")
        for filename in tarHandle.getnames():
            tarHandle.extract(filename, desdir)
            if filename.endswith('.tar.gz'):
                path = desdir
                zfile = filename
                self.untar_dir(path, zfile)
        tarHandle.close()

    # 解压文件
    def un_zip(self, file_path):
        (filepath, file_name) = os.path.split(file_path)
        (filename_zip, extension_zip) = os.path.splitext(file_name)
        # 切换指定存储的文件夹
        os.chdir(filepath)
        if extension_zip == '.zip':
            try:
                self.unzip(filepath, file_name)
            except Exception as e:
                return (False, e, extension_zip)
        elif extension_zip == '.gz':
            try:
                self.untar_dir(filepath, file_name)
            except Exception as e:
                return (False, e, extension_zip)

    # 遍历scripts下的文件夹,给pdf加上水印
    def dir_scripts(self, rootdir):
        for list in os.listdir(rootdir):
            path = os.path.join(rootdir, list)
            if '.pdf' in path:
                pdf_file_in = path
                pdf_file_out = path
                pdf_file_mark = self.create_watermark(user_name)
                self.add_watermark(pdf_file_in, pdf_file_mark, pdf_file_out)
            if os.path.isdir(path):
                self.dir_scripts(path)

    # 在 scripts目录下进行压缩zip
    def zip_file_path(self, file_name):
        file_name = self.six_code + file_name
        fantasy_zip = zipfile.ZipFile(os.path.join(os.getcwd(), file_name), 'w')
        for folder, subfolders, files in os.walk(os.getcwd()):
            for file in files:
                if '.zip' not in file:
                    fantasy_zip.write(os.path.join(folder, file), file, compress_type=zipfile.ZIP_DEFLATED)
        fantasy_zip.close()

    # 在scripts目录下进形压缩tar.gz
    def tar_gz_file_path(self, file_name):
        # 创建压缩包名
        file_name = self.six_code + file_name
        tar = tarfile.open(os.path.join(os.getcwd(), file_name), 'w:gz')
        # 创建压缩包
        for root, dir, files in os.walk('./'):
            for file in files:
                fullpath = os.path.join(root, file)
                tar.add(fullpath)
        tar.close()

    # 生成6位数字的字符串
    def six_random(self):
        six_str = ""
        for i in range(6):
            ch = chr(random.randrange(ord('0'), ord('9') + 1))
            six_str += ch
        return six_str
path = '路径'
user_name = 'wangfe'
water_mark = watermark_sv(path, user_name)