<filename>tests/test_publickey.py

import os
import tempfile
import base64
import unittest
try:
    from unitttest import mock
except ImportError:
    import mock

import Crypto.PublicKey.RSA
import Crypto.Cipher.PKCS1_v1_5

from edenred.publickey import PublicKey

PRIVATE_KEY = """\
-----<KEY>
"""


class TestPublicKey(unittest.TestCase):
    def __init__(self, *args, **kwargs):
        self.private = Crypto.PublicKey.RSA.importKey(PRIVATE_KEY)
        self.public = self.private.publickey()
        super(TestPublicKey, self).__init__(*args, **kwargs)

    def setUp(self):
        _, self.private_path = tempfile.mkstemp(suffix='.pem', prefix='private-')
        with open(self.private_path, 'w') as private_key_file:
            private_key_file.write(self.private.exportKey('PEM'))

        _, self.public_path = tempfile.mkstemp(suffix='.pem', prefix='public-')
        with open(self.public_path, 'w') as public_key_file:
            public_key_file.write(self.public.exportKey('PEM'))

    def tearDown(self):
        os.unlink(self.private_path)
        os.unlink(self.public_path)

    def assert_pkcs1_b64_encrypted(self, message, encrypted):
        encrypted = base64.b64decode(encrypted.encode())
        cipher = Crypto.Cipher.PKCS1_v1_5.new(self.private)
        sentinel = Crypto.Random.new().read(len(message))

        decrypted = cipher.decrypt(encrypted, sentinel)

        self.assertEqual(message, decrypted)

    def test_init(self):
        key = PublicKey(self.public_path)

        self.assertEqual(self.public, key.rsa)
        self.assertFalse(key.testing)

    def test_init_testing(self):
        testing = mock.Mock(spec=bool)

        key = PublicKey(self.public_path, testing)

        self.assertEqual(self.public, key.rsa)
        self.assertEqual(testing, key.testing)

    def test_encrypt_testing(self):
        key = PublicKey(self.public_path, testing=True)

        data = mock.Mock(spec=str)

        self.assertEqual(data, key.encrypt(data))

    def test_encrypt(self):
        key = PublicKey(self.public_path)
        message = "123456789abcdefghij"

        encrypted = key.encrypt(message)

        self.assert_pkcs1_b64_encrypted(message, encrypted)

    def test_equal(self):
        testing = mock.Mock(spec=bool)
        key1 = PublicKey(self.public_path, testing)
        key2 = PublicKey(self.public_path, testing)

        self.assertEqual(key1, key2)
