Minor improvements to CardGenerator and the corresponding unit test

Make sure that the file gets closed when writing a card to/reading a
card from disk. Also add a unit test for the error condition when trying
to import a card from a nonexistent file.
This commit is contained in:
Dominik
2014-12-21 16:09:26 +01:00
parent 1e7492a806
commit 991ba148d9
2 changed files with 24 additions and 15 deletions

View File

@@ -382,16 +382,19 @@ class CardGenerator(object):
def loadCard(self, filename): def loadCard(self, filename):
"""Load a card from disk""" """Load a card from disk"""
db = anydbm.open(filename, 'r')
if self.password is None: if self.password is None:
self.password = getpass.getpass("Please enter your password:") self.password = getpass.getpass("Please enter your password:")
serializedMF = read_protected_string(db["mf"], self.password) db = anydbm.open(filename, 'r')
serializedSAM = read_protected_string(db["sam"], self.password) try:
serializedMF = read_protected_string(db["mf"], self.password)
serializedSAM = read_protected_string(db["sam"], self.password)
self.type = db["type"]
finally:
db.close()
self.sam = loads(serializedSAM) self.sam = loads(serializedSAM)
self.mf = loads(serializedMF) self.mf = loads(serializedMF)
self.type = db["type"]
def saveCard(self, filename): def saveCard(self, filename):
"""Save the currently running card to disk""" """Save the currently running card to disk"""
@@ -413,11 +416,13 @@ class CardGenerator(object):
protectedSAM = protect_string(sam_string, self.password) protectedSAM = protect_string(sam_string, self.password)
db = anydbm.open(filename, 'c') db = anydbm.open(filename, 'c')
db["mf"] = protectedMF try:
db["sam"] = protectedSAM db["mf"] = protectedMF
db["type"] = self.type db["sam"] = protectedSAM
db["version"] = "0.1" db["type"] = self.type
db.close() db["version"] = "0.1"
finally:
db.close()
if __name__ == "__main__": if __name__ == "__main__":
from optparse import OptionParser from optparse import OptionParser

View File

@@ -17,9 +17,10 @@
# virtualsmartcard. If not, see <http://www.gnu.org/licenses/>. # virtualsmartcard. If not, see <http://www.gnu.org/licenses/>.
# #
import unittest import anydbm
import tempfile
import os import os
import tempfile
import unittest
from virtualsmartcard.CardGenerator import CardGenerator from virtualsmartcard.CardGenerator import CardGenerator
class TestNPACardGenerator(unittest.TestCase): class TestNPACardGenerator(unittest.TestCase):
@@ -29,15 +30,13 @@ class TestNPACardGenerator(unittest.TestCase):
self.nPA_generator = CardGenerator('nPA') self.nPA_generator = CardGenerator('nPA')
self.nPA_generator.password = "TestPassword" self.nPA_generator.password = "TestPassword"
def tearDown(self):
os.unlink(self.filename)
def test_nPA_creation(self): def test_nPA_creation(self):
self.nPA_generator.generateCard() self.nPA_generator.generateCard()
self.nPA_generator.saveCard(self.filename) self.nPA_generator.saveCard(self.filename)
mf, sam = self.nPA_generator.getCard() mf, sam = self.nPA_generator.getCard()
self.assertIsNotNone(mf) self.assertIsNotNone(mf)
self.assertIsNotNone(sam) self.assertIsNotNone(sam)
os.unlink(self.filename)
def test_load_nPA_from_file_nPA_from_file(self): def test_load_nPA_from_file_nPA_from_file(self):
self.nPA_generator.generateCard() self.nPA_generator.generateCard()
@@ -48,6 +47,11 @@ class TestNPACardGenerator(unittest.TestCase):
mf, sam = local_generator.getCard() mf, sam = local_generator.getCard()
self.assertIsNotNone(mf) self.assertIsNotNone(mf)
self.assertIsNotNone(sam) self.assertIsNotNone(sam)
os.unlink(self.filename)
def test_load_nonexistent_file(self):
with self.assertRaises(anydbm.error):
self.nPA_generator.loadCard(self.filename)
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()