# Cross-Platform Identity Authentication and End-to-End Encrypted Communication Technology Based on DID
# Author: GaoWei Chang
# Email: chgaowei@gmail.com
# Website: http://pi-unlimited.com
#
# This project is open-sourced under the MIT License. For details, please see the LICENSE file.
# Version: 0.1.0

import unittest
import asyncio
import json
from typing import Dict, Any
from cryptography.hazmat.primitives import serialization
import sys
import os
import logging

current_script_path = os.path.abspath(__file__)
current_directory = os.path.dirname(current_script_path)
sys.path.append(current_directory)
sys.path.append(current_directory + "/../")
sys.path.append(current_directory + "/../../")

from didall.short_term_key_generater import ShortTermKeyGenerater
from didall.utils.did_generate import did_generate

class TestShortTermKeyGeneration(unittest.TestCase):
    def setUp(self):
        # 使用did_generate生成发起方和接收方的DID和私钥
        initiator_private_key, initiator_public_key, initiator_did, _ = did_generate("wss://initiator.com/endpoint")
        receiver_private_key, receiver_public_key, receiver_did, _ = did_generate("wss://receiver.com/endpoint")

        # 将私钥保存为PEM格式
        initiator_private_key_pem = initiator_private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        receiver_private_key_pem = receiver_private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )

        initiator_private_key_pem = initiator_private_key_pem.decode('utf-8')
        receiver_private_key_pem = receiver_private_key_pem.decode('utf-8')

        # 初始化两个实例，一个作为发起方，一个作为接收方
        self.initiator = ShortTermKeyGenerater(
            local_did=initiator_did,
            did_private_key_pem=initiator_private_key_pem,
            remote_did=receiver_did,
            json_send_func=self.receiver_receive_json,
            is_initiator=True
        )

        self.receiver = ShortTermKeyGenerater(
            local_did=receiver_did,
            did_private_key_pem=receiver_private_key_pem,
            remote_did=initiator_did,
            json_send_func=self.initiator_receive_json,
            is_initiator=False,
            session_id=self.initiator.session_id
        )

        self.initiator_did = initiator_did
        self.receiver_did = receiver_did

    async def initiator_receive_json(self, json_data: Dict[str, Any]):
        # 等待10ms
        self.initiator.receive_json_message(json_data)
        await asyncio.sleep(0.01)

    async def receiver_receive_json(self, json_data: Dict[str, Any]):
        self.receiver.receive_json_message(json_data)
        await asyncio.sleep(0.01)

    async def initiate_key_exchange(self):
        result = await self.initiator.generate_short_term_key_active()
        logging.info(f'generate_short_term_key_active {result}')
        self.assertTrue(result, "密钥交换失败")
    
    async def receive_key_exchange(self):
        """启动密钥交换过程"""
        result = await self.receiver.generate_short_term_key_passive()
        logging.info(f'generate_short_term_key_passive {result}')
        self.assertTrue(result, "密钥交换失败")


    def test_key_exchange(self):
        # 测试密钥交换过程
        async def run_test():
            return await asyncio.gather(
                self.initiate_key_exchange(),
                self.receive_key_exchange()
            )

        asyncio.run(run_test())

        # 验证双方是否协商出相同的密钥
        initiator_remote_did, \
        initiator_send_encryption_key, \
        initiator_receive_decryption_key, \
        initiator_secret_key_id, \
        initiator_key_expires, \
        initiator_cipher_suite = self.initiator.get_final_short_term_key()

        receiver_remote_did, \
        receiver_send_encryption_key, \
        receiver_receive_decryption_key, \
        receiver_secret_key_id, \
        receiver_key_expires, \
        receiver_cipher_suite = self.receiver.get_final_short_term_key()

        self.assertEqual(initiator_remote_did, self.receiver_did, "远程DID不匹配")
        self.assertEqual(receiver_remote_did, self.initiator_did, "远程DID不匹配")
        self.assertEqual(initiator_send_encryption_key, receiver_receive_decryption_key, "发送加密密钥不匹配")
        self.assertEqual(initiator_receive_decryption_key, receiver_send_encryption_key, "接收解密密钥不匹配")
        self.assertEqual(initiator_secret_key_id, receiver_secret_key_id, "密钥ID不匹配")
        self.assertEqual(initiator_key_expires, receiver_key_expires, "密钥有效期不匹配")
        self.assertEqual(initiator_cipher_suite, receiver_cipher_suite, "密码套件不匹配")

        # logging.info(f"initiator_remote_did: {initiator_remote_did}")
        # logging.info(f"receiver_remote_did: {receiver_remote_did}")
        # logging.info(f"initiator_send_encryption_key: {initiator_send_encryption_key}")
        # logging.info(f"initiator_receive_decryption_key: {initiator_receive_decryption_key}")
        # logging.info(f"initiator_secret_key_id: {initiator_secret_key_id}")
        # logging.info(f"receiver_send_encryption_key: {receiver_send_encryption_key}")
        # logging.info(f"receiver_receive_decryption_key: {receiver_receive_decryption_key}")
        # logging.info(f"receiver_secret_key_id: {receiver_secret_key_id}")
        # logging.info(f"initiator_key_expires: {initiator_key_expires}")
        # logging.info(f"receiver_key_expires: {receiver_key_expires}")
        # logging.info(f"initiator_cipher_suite: {initiator_cipher_suite}")
        # logging.info(f"receiver_cipher_suite: {receiver_cipher_suite}")

if __name__ == "__main__":
    unittest.main()

