test_cert_service

import unittest from unittest.mock import patch, MagicMock, mock_open from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.primitives.serialization import pkcs12 from cert_service import CertService from Crypto.Cipher import AES from Crypto.Util.Padding import pad, unpad import base64 class TestCertService(unittest.TestCase): def setUp(self): self.config = { "certs": { "test_category": { "CERT_PATH": "mock_path", "CERT_FILE": "mock_file.p12", "CERT_PASSWORD": "mock_password", "AES_ENCRYPTION_KEY": base64.b64encode(b"mock_aes_key" * 2).decode() # 32 bytes } } } self.cert_category = "test_category" @patch("builtins.open", new_callable=mock_open, read_data=b"mock_cert_data") @patch("cryptography.hazmat.primitives.serialization.pkcs12.load_key_and_certificates") @patch("base64.b64decode") @patch("Crypto.Cipher.AES.new") def test_init_with_aes_key(self, mock_aes_new, mock_b64decode, mock_load_cert, mock_open_file): mock_private_key = MagicMock() mock_certificate = MagicMock() mock_load_cert.return_value = (mock_private_key, mock_certificate, None) mock_b64decode.side_effect = [b"mock_aes_key", b"mock_aes_key"] # AES key decoded twice service = CertService(self.config, self.cert_category) self.assertEqual(service.private_key, mock_private_key) self.assertEqual(service.public_key, mock_certificate.public_key()) mock_aes_new.assert_called_once_with(b"mock_aes_key", AES.MODE_ECB) @patch("builtins.open", new_callable=mock_open, read_data=b"mock_cert_data") @patch("cryptography.hazmat.primitives.serialization.pkcs12.load_key_and_certificates") @patch("base64.b64decode") def test_init_without_aes_key(self, mock_b64decode, mock_load_cert, mock_open_file): del self.config["certs"][self.cert_category]["AES_ENCRYPTION_KEY"] mock_private_key = MagicMock() mock_certificate = MagicMock() mock_load_cert.return_value = (mock_private_key, mock_certificate, None) service = CertService(self.config, self.cert_category) self.assertEqual(service.private_key, mock_private_key) self.assertEqual(service.public_key, mock_certificate.public_key()) self.assertFalse(hasattr(service, "cipher")) def test_encrypt(self): service = MagicMock(spec=CertService) service.public_key.encrypt.return_value = b"mock_encrypted_data" encrypted_data = service.encrypt(b"mock_data") service.public_key.encrypt.assert_called_once_with( b"mock_data", padding=pkcs12.OAEP( mgf=pkcs12.MGF1(algorithm=pkcs12.SHA256()), algorithm=pkcs12.SHA256(), label=None ) ) self.assertEqual(encrypted_data, b"mock_encrypted_data") def test_decrypt(self): service = MagicMock(spec=CertService) service.private_key.decrypt.return_value = b"mock_decrypted_data" decrypted_data = service.decrypt(b"mock_encrypted_data") service.private_key.decrypt.assert_called_once_with( b"mock_encrypted_data", padding=pkcs12.OAEP( mgf=pkcs12.MGF1(algorithm=pkcs12.SHA256()), algorithm=pkcs12.SHA256(), label=None ) ) self.assertEqual(decrypted_data, b"mock_decrypted_data") @patch("Crypto.Cipher.AES.new") def test_encrypt_aes_with_custom_key(self, mock_aes_new): mock_cipher = MagicMock() mock_aes_new.return_value = mock_cipher mock_cipher.encrypt.return_value = b"mock_encrypted_data" service = CertService(self.config, self.cert_category) key = b"mock_custom_key" * 2 # 32 bytes data = b"mock_data" encrypted_data = service.encrypt_aes(data, key=key) mock_aes_new.assert_called_once_with(key, AES.MODE_ECB) mock_cipher.encrypt.assert_called_once_with(pad(data, AES.block_size)) self.assertEqual(encrypted_data, b"mock_encrypted_data") def test_encrypt_aes_with_default_key(self): mock_cipher = MagicMock() service = CertService(self.config, self.cert_category) service.cipher = mock_cipher mock_cipher.encrypt.return_value = b"mock_encrypted_data" data = b"mock_data" encrypted_data = service.encrypt_aes(data) mock_cipher.encrypt.assert_called_once_with(pad(data, AES.block_size)) self.assertEqual(encrypted_data, b"mock_encrypted_data") @patch("Crypto.Cipher.AES.new") def test_decrypt_aes_with_custom_key(self, mock_aes_new): mock_cipher = MagicMock() mock_aes_new.return_value = mock_cipher mock_cipher.decrypt.return_value = pad(b"mock_data", AES.block_size) service = CertService(self.config, self.cert_category) key = b"mock_custom_key" * 2 # 32 bytes encrypted_data = b"mock_encrypted_data" decrypted_data = service.decrypt_aes(encrypted_data, key=key) mock_aes_new.assert_called_once_with(key, AES.MODE_ECB) mock_cipher.decrypt.assert_called_once_with(encrypted_data) self.assertEqual(decrypted_data, b"mock_data") def test_decrypt_aes_with_default_key(self): mock_cipher = MagicMock() service = CertService(self.config, self.cert_category) service.cipher = mock_cipher mock_cipher.decrypt.return_value = pad(b"mock_data", AES.block_size) encrypted_data = b"mock_encrypted_data" decrypted_data = service.decrypt_aes(encrypted_data) mock_cipher.decrypt.assert_called_once_with(encrypted_data) self.assertEqual(decrypted_data, b"mock_data") if __name__ == "__main__": unittest.main()

Posted By