import sys
import os
from PyQt6.QtWidgets import (QApplication, QWidget, QPushButton, QVBoxLayout, 
                            QHBoxLayout, QFileDialog, QLineEdit, QLabel, QMessageBox,
                            QProgressBar, QTextEdit, QSplitter)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QFont, QColor
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
 
class CryptoWorker(QThread):
    progress = pyqtSignal(int)
    finished = pyqtSignal()
    error = pyqtSignal(str)
    status = pyqtSignal(str)  # 新增狀態(tài)信號
 
    def __init__(self, folder, password, mode):
        super().__init__()
        self.folder = folder
        self.password = password
        self.mode = mode  # 'encrypt' or 'decrypt'
 
    def run(self):
        try:
            key = self.get_key(self.password)
            fernet = Fernet(key)
 
            total_files = sum([len(files) for r, d, files in os.walk(self.folder)])
            processed_files = 0
            
            self.status.emit(f"開始{'加密' if self.mode == 'encrypt' else '解密'}操作...\n")
 
            for root, dirs, files in os.walk(self.folder):
                for file in files:
                    file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(file_path, self.folder)
                    
                    if self.mode == 'encrypt':
                        if file.endswith('.enc'):
                            self.status.emit(f"跳過已加密文件: {relative_path}")
                        else:
                            self.encrypt_file(fernet, file_path)
                            self.status.emit(f"已加密: {relative_path}")
                    else:  # decrypt
                        if not file.endswith('.enc'):
                            self.status.emit(f"跳過未加密文件: {relative_path}")
                        else:
                            self.decrypt_file(fernet, file_path)
                            self.status.emit(f"已解密: {relative_path}")
                    
                    processed_files += 1
                    self.progress.emit(int(processed_files / total_files * 100))
 
            self.status.emit("\n操作完成！")
            self.finished.emit()
        except Exception as e:
            self.error.emit(str(e))
 
    def get_key(self, password):
        password = password.encode()
        salt = b'salt_'
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password))
        return key
 
    def encrypt_file(self, fernet, file_path):
        with open(file_path, 'rb') as f:
            data = f.read()
        encrypted = fernet.encrypt(data)
        encrypted_path = file_path + '.enc'
        with open(encrypted_path, 'wb') as f:
            f.write(encrypted)
        os.remove(file_path)
 
    def decrypt_file(self, fernet, file_path):
        with open(file_path, 'rb') as f:
            data = f.read()
        try:
            decrypted = fernet.decrypt(data)
            decrypted_path = file_path[:-4]  # 移除 '.enc'
            with open(decrypted_path, 'wb') as f:
                f.write(decrypted)
            os.remove(file_path)
        except:
            raise Exception(f'解密 {file_path} 失敗，可能密碼錯誤')
 
class Encryptor(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
 
    def initUI(self):
        self.setWindowTitle('文件夾加密器')
        self.setGeometry(300, 200, 700, 600)
 
        main_layout = QVBoxLayout()
 
        # 上半部分：控制面板
        control_panel = QWidget()
        control_layout = QVBoxLayout()
 
        # 文件夾選擇
        folder_layout = QHBoxLayout()
        self.folder_input = QLineEdit()
        folder_button = QPushButton('選擇文件夾')
        folder_button.clicked.connect(self.select_folder)
        folder_layout.addWidget(self.folder_input)
        folder_layout.addWidget(folder_button)
        control_layout.addLayout(folder_layout)
 
        # 密碼輸入
        password_layout = QHBoxLayout()
        password_label = QLabel('密碼:')
        self.password_input = QLineEdit()
        self.password_input.setEchoMode(QLineEdit.EchoMode.Password)
        password_layout.addWidget(password_label)
        password_layout.addWidget(self.password_input)
        control_layout.addLayout(password_layout)
 
        # 進(jìn)度條
        self.progress_bar = QProgressBar()
        control_layout.addWidget(self.progress_bar)
 
        # 加密解密按鈕
        button_layout = QHBoxLayout()
        encrypt_button = QPushButton('加密')
        decrypt_button = QPushButton('解密')
        encrypt_button.clicked.connect(lambda: self.process_folder('encrypt'))
        decrypt_button.clicked.connect(lambda: self.process_folder('decrypt'))
        button_layout.addWidget(encrypt_button)
        button_layout.addWidget(decrypt_button)
        control_layout.addLayout(button_layout)
 
        control_panel.setLayout(control_layout)
        main_layout.addWidget(control_panel)
 
        # 分割線
        splitter = QSplitter(Qt.Orientation.Vertical)
 
        # 注意事項面板
        notice_panel = QWidget()
        notice_layout = QVBoxLayout()
        notice_label = QLabel("注意事項:")
        notice_label.setFont(QFont("Arial", 12, QFont.Weight.Bold))
        notice_label.setStyleSheet("color: red;")
        notice_text = QTextEdit()
        notice_text.setPlainText(
            "1. 請確保記住密碼，否則將無法解密文件。\n"
            "2. 請慎重操作！該程序會直接在原文件夾中操作，謹(jǐn)防誤操作。\n"
            "3. 加密后刪除原文件，解密后刪除加密文件。加密后的文件會有 .enc 擴展名。\n"
            "4. 本程序支持加密/解密文件夾中的所有文件。\n"
            "5. 已加密的文件不會被重復(fù)加密，未加密的文件不會被解密。"
        )
        notice_text.setReadOnly(True)
        notice_text.setFont(QFont("Arial", 11))
        notice_text.setStyleSheet("color: #D32F2F;")  # 深紅色
        notice_layout.addWidget(notice_label)
        notice_layout.addWidget(notice_text)
        notice_panel.setLayout(notice_layout)
 
        # 日志面板
        log_panel = QWidget()
        log_layout = QVBoxLayout()
        log_label = QLabel("處理日志:")
        #log_label.setFont(QFont("Arial", 10, QFont.Weight.Bold))
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        #self.log_text.setFont(QFont("Arial", 10))
        log_layout.addWidget(log_label)
        log_layout.addWidget(self.log_text)
        log_panel.setLayout(log_layout)
 
        # 添加面板到分割器
        splitter.addWidget(notice_panel)
        splitter.addWidget(log_panel)
        main_layout.addWidget(splitter)
 
        self.setLayout(main_layout)
 
    def select_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "選擇文件夾")
        self.folder_input.setText(folder)
 
    def process_folder(self, mode):
        folder = self.folder_input.text()
        password = self.password_input.text()
        if not folder or not password:
            QMessageBox.warning(self, '警告', '請選擇文件夾并輸入密碼')
            return
 
        # 清空日志
        self.log_text.clear()
        
        self.worker = CryptoWorker(folder, password, mode)
        self.worker.progress.connect(self.update_progress)
        self.worker.finished.connect(self.process_finished)
        self.worker.error.connect(self.process_error)
        self.worker.status.connect(self.update_log)  # 連接狀態(tài)信號
        self.worker.start()
 
    def update_progress(self, value):
        self.progress_bar.setValue(value)
 
    def update_log(self, message):
        self.log_text.append(message)
 
    def process_finished(self):
        QMessageBox.information(self, '成功', '操作完成')
        self.progress_bar.setValue(0)
 
    def process_error(self, error_msg):
        QMessageBox.critical(self, '錯誤', f'發(fā)生錯誤：{error_msg}')
        self.progress_bar.setValue(0)
        self.log_text.append(f"\n錯誤: {error_msg}")
 
if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = Encryptor()
    ex.show()
    sys.exit(app.exec())