# -*- coding: utf-8 -*-
# ☯ Author: ChinaPython
# ☯ Date  : 2021/9/16 19:01
import os
import PySimpleGUI as sg
from PyPDF2 import PdfFileWriter
from PyPDF2 import PdfFileReader


class EncryptTool:
    def __init__(self):
        self.layout = [
            [
                sg.Text(text="阅览密码"),
                sg.InputText(do_not_clear=True, default_text="默认:admin", text_color='gray', key='psw_a', size=(17, 1)),
                sg.Text(text="编辑密码"),
                sg.InputText(do_not_clear=True, default_text="默认:admin", text_color='gray', key='psw_b', size=(17, 1)),
            ],
            [
                sg.Input(do_not_clear=True, default_text='选择需要转换的PDF文件(默认同路径)', text_color='gray', key='path'),
                sg.FileBrowse(button_text='选择文件', key='file'),
            ],
            [
                sg.Input(do_not_clear=True, default_text='转换后PDF文件存储路径', text_color='gray', key='path2'),
                sg.FolderBrowse(button_text='输出目录', key='folder'),
            ],
            [sg.Text(text="", size=(48, 4), background_color="white", text_color="red", key="msg")],
            [sg.Button(button_text='开始加密', key='start', size=(48, 2))],
        ]
        self.window = sg.Window('PDF加密小工具(吾爱破解论坛)', self.layout, icon='icon.ico')

    @staticmethod
    def __encrypt(unencrypted_pdf, read_password, encrypted_pdf=None, owner_password=None):
        unencrypted_pdf = "加密_" + unencrypted_pdf if encrypted_pdf is None else unencrypted_pdf
        owner_password = owner_password if owner_password is None else owner_password
        try:
            content = PdfFileReader(unencrypted_pdf)
            writer = PdfFileWriter()
            [writer.addPage(content.getPage(page)) for page in range(content.getNumPages())]
            writer.encrypt(user_pwd=read_password, owner_pwd=owner_password, use_128bit=False)
            with open(encrypted_pdf, 'wb') as out:
                writer.write(out)
            return os.path.exists(encrypted_pdf)
        except Exception as e:
            return e

    # 运行入口
    def run(self):
        while True:
            event, values = self.window.Read()
            if event == 'start':
                a, b, c, d = values['psw_a'], values['psw_b'], values['path'], values['path2']
                if os.path.exists(c) is False or str(c).lower().endswith(".pdf") is False:
                    self.window.Element('msg').Update("加载目录失败，请检测路径的合法性")
                    continue
                a = "admin" if "默认:" in a or len(a) == 0 else a
                b = "admin" if "默认:" in b or len(b) == 0 else b
                filepath, filename = os.path.split(c)
                d = f"{d}/加密_{filename}" if os.path.isdir(d) is True else f"{filepath}/加密_{filename}"
                msg = self.__encrypt(c, a, d, b)
                if msg is not True:
                    self.window.Element('msg').Update(f"加密失败:{msg}")
                    continue
                self.window.Element('msg').Update(
                    f"阅览密码:{a}\n编辑密码:{b}\n目标文件:{c}\n输出文件:{d}"
                )
            if event is None:
                break
        self.window.close()


if __name__ == '__main__':
    start = EncryptTool()
    start.run()
