from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric import utils
import os

def verify_digital_signature():
    # 获取当前脚本所在目录
    current_directory = os.path.dirname(__file__)

    # 读取A的公钥
    public_key_path = os.path.join(current_directory, 'A_public.txt')
    with open(public_key_path, 'r') as file:
        public_key_data = file.read()
        public_key = serialization.load_pem_public_key(
            public_key_data.encode(),
            backend=default_backend()
        )

    # 读取解密后的文件
    decrypted_file_path = os.path.join(current_directory, 'file.txt_decrypted.txt')
    with open(decrypted_file_path, 'rb') as file:
        decrypted_message = file.read()

    # 计算摘要
    hash_object = hashes.Hash(hashes.SHA256(), backend=default_backend())
    hash_object.update(decrypted_message)
    hash_value = hash_object.finalize()

    # 读取数字签名
    signature_path = os.path.join(current_directory, 'signature.txt')
    with open(signature_path, 'rb') as file:
        signature = file.read()

    # 使用公钥验证签名
    try:
        public_key.verify(
            signature,
            hash_value,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            utils.Prehashed(hashes.SHA256())
        )
        print("数字签名验证成功！")
    except Exception as e:
        print(f"数字签名验证失败: {e}")