import os
import unittest

from BaDCA import Keys, CSRs

class BaDCAKeysTestCase(unittest.TestCase):
    obj = None

    def subjectCheck(self, which, val):
        ckVal = self.obj.getSubject(which)
        assert ckVal == val, "Subject '%s' failed! Returned %s instead of %s" \
                       % (which, str(ckVal), val)

    def subjectSet(self, which, val):
        assert self.obj.setSubject(which, val) == 1, \
               "Failed to set subject '%s' to value '%s'" % (which, val)

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

    def test01Read(self):
        """ Test reading of a CSR from a file """
        assert self.obj.readFromFile('tests/csr/test1.csr') == 1, \
                                                "Failed to read the CSR"

    def test02Read2(self):
        """ Test reading of a CSR from a string """
        f = open('tests/csr/test1.csr', 'r')
        txt = f.read()
        f.close()
        assert self.obj.readFromMemory(txt) == 1, \
                                     "Failed to parse the CSR in memory"

    def test03Key(self):
        """ Test key extraction from a CSR """
        assert self.obj.readFromFile('tests/csr/test1.csr') == 1, \
                                                "Failed to read the CSR"
        key = self.obj.getKey()
        assert key is not None, "Unable to get Key object from CSR"
        assert key.hasPublic(), "No public key found"
        assert key.hasPrivate() == 0, "Private key found when none should exist"
        assert key.bits == 2048, "Incorrect strength key returned"

    def test04Info(self):
        """ Test extraction of information from a CSR """
        assert self.obj.readFromFile('tests/csr/test1.csr') == 1, \
                                                "Failed to read the CSR"
        self.subjectCheck('C', 'GB')
        self.subjectCheck('CN', "asylum.zones.apache.org")
        self.subjectCheck('Email', "dreid@apache.org")

    def test05(self):
        """ CSR with a private key available """
        assert self.obj.readFromFile('tests/csr/test2.csr') == 1, \
                                                "Failed to read the CSR"
        key = self.obj.getKey()
        assert key is not None, "Unable to get key object from CSR"
        assert key.hasPublic(), "No public key found!"
        assert key.addSearchDirectory('tests/keys/private') == 1, \
                                 "Unable to add search directory to key"
        assert key.hasPrivate() == 1, "Unable to find private key"
        assert key.hasPublic(), "No public key found!"
        self.subjectCheck('C', 'GB')
        self.subjectCheck('CN', "david reid")
        self.subjectCheck('Email', "dreid@apache.org")


    def test06(self):
        """ Create a CSR """
        key = Keys.RSAKey()
        key.addSearchDirectory('tests/keys/private')
        key.readPublicKey('tests/keys/public/test3.public.key')
        assert key.hasPrivate(), "Failed to get private key"
        self.subjectSet('C', 'GB')
        self.subjectSet('CN', 'test06')
        self.subjectSet('Email', 'test@test.com')
        self.obj.setKey(key)
        assert self.obj.createCSR() == 1, "Error creating CSR"
        assert self.obj.writeToDirectory('tests/csr') == 1, \
                                "Unable to write to the temporary file!"
        fn = self.obj.getFilename()
        assert os.path.exists(fn), "File was not created"
        os.unlink(fn)
        assert os.path.exists(fn) == 0, "File was not removed cleanly"

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