import unittest

from BaDCA import Keys
from BaDCA.Testing import *

class BaDCAKeysTestCase(unittest.TestCase):
    obj = None

    def setUp(self):
        """ Called prior to every test """
        if self.obj is None:
            self.obj = Keys.RSAKey()
        else:
            self.Reset()

#    def test01(self):
#        """ Testing creation of keys of different strengths """
#        for r in ( 1024, 2048, 4096 ):
#            assert self.obj.create(r) == 1, \
#                            "Unable to create key of strength " + str(r)
#            assert self.obj.bits == r, "Incorrect strength key created"

    def test02(self):
        """ Checking we accept only keys of correct strengths """
        self.assertRaises(Keys.rsa.keySize, self.obj.create, 512)
        self.assertRaises(Keys.rsa.keySize, self.obj.create, 9216)
        self.assertRaises(Keys.rsa.keySize, self.obj.create, 1023)
        self.assertRaises(Keys.rsa.keySize, self.obj.create, 8193)

    def test03(self):
        """ Test reading in a public key """
        assert self.obj.readPublicKey('tests/keys/test1.public.key') == 1, \
                                             "Failed to read public key"
        assert self.obj.isValid(), "Invalid key object created"
        assert self.obj.hasPublic(), "No public key available"

    def test04(self):
        """ Test reading of private key """
        assert self.obj.readPrivateKey('tests/keys/test2.key') == 1, \
                                            "Failed to read private key"
        assert self.obj.isValid(), "Invalid key object created"
        assert self.obj.hasPublic(), "Public key not available??"
        assert self.obj.hasPrivate(), "Private key not available??"

    def test06(self):
        """ Test directory support """
        assert self.obj.readPublicKey('tests/keys/public/test3.public.key') == 1, \
                                             "Failed to read public key"
        assert self.obj.hasPublic(), "No public key available"
        assert self.obj.hasPrivate() == 0, "Private key available???"
        assert self.obj.addSearchDirectory('tests/keys/private') == 1, \
                                 "Unable to add search directory to key"
        assert self.obj.hasPrivate(), "Private key not available"

    def test07(self):
        """ Test extracting a key from a certificate """
        certObj = readCertificate('tests/certs/test1.pem')
        assert certObj is not None, "Failed to read certificate!"
        self.obj = Keys.RSAKey(certificate = certObj)
        assert self.obj.isValid(), "Invalid object created from certificate"
        assert self.obj.hasPublic(), "Failed to process key from certificate"
        assert self.obj.hasPrivate() is False, \
                     "The private key is available?"

    def test08(self):
        """ Test extracting a key from a CSR """
        csrObj = readCSR('tests/csr/test1.csr')
        assert csrObj is not None, "Failed to read CSR!"
        self.obj = Keys.RSAKey(csr = csrObj)
        assert self.obj.isValid(), "Invalid object created from CSR"
        assert self.obj.hasPublic(), "Failed to process key from CSR"
        assert self.obj.hasPrivate() is False, \
                     "The private key is available?"

    def test09(self):
        """ Test reading from memory """
        f = open('tests/keys/public/test3.public.key', "r")
        txt = f.read()
        f.close()
        assert len(txt) > 0, "Failed to read text!"
        assert self.obj.readPublicKeyFromMemory(txt) == 1, \
                                      "Failed to create key from memory"
        assert self.obj.hasPublic(), "Error creating key"

    def test10(self):
        """ Test key equality checks """
        k1 = Keys.RSAKey(filename = 'tests/keys/test1.public.key')
        k2 = Keys.RSAKey(filename = 'tests/keys/test2.public.key')
        assert k1 != k2, "The keys should NOT be equal"
        assert k1 == k1, "Same key should equal itself"

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