瀏覽代碼

Implement save and load key for mongodb

oz123 10 年之前
父節點
當前提交
4345c3e481
共有 2 個文件被更改,包括 12 次插入7 次删除
  1. 8 2
      pwman/data/drivers/mongodb.py
  2. 4 5
      tests/test_mongodb.py

+ 8 - 2
pwman/data/drivers/mongodb.py

@@ -32,6 +32,7 @@ class MongoDB(Database):
 
     def _open(self):
         self._con = pymongo.Connection(self.uri)
+        self._db = self._con.get_default_database()
 
     def getnodes(self, ids):
         pass
@@ -55,10 +56,15 @@ class MongoDB(Database):
         pass
 
     def savekey(self, key):
-        pass
+        coll = self._db['crypto']
+        salt, digest = key.split('$6$')
+        coll.insert({'salt': salt, 'key': digest})
 
     def loadkey(self):
-        pass
+        coll = self._db['crypto']
+        key = coll.find_one({}, {'_id': 0})
+        key = key['salt'] + '$6$' + key['key']
+        return key
 
     def close(self):
         pass

+ 4 - 5
tests/test_mongodb.py

@@ -49,21 +49,20 @@ class TestMongoDB(unittest.TestCase):
 
     @classmethod
     def tearDownClass(cls):
-        # Drop collections here
-        pass
+        coll = cls.db._db['crypto']
+        coll.drop()
 
     def test_1_con(self):
         self.assertIsInstance(self.db._con, pymongo.Connection)
 
-    @unittest.skip("")
+    @unittest.skip("MongoDB creates collections on the fly")
     def test_2_create_collections(self):
         pass
 
-    @unittest.skip("")
     def test_3_load_key(self):
         self.db.savekey('SECRET$6$KEY')
         secretkey = self.db.loadkey()
-        self.assertEqual(secretkey, 'SECRET$6$KEY')
+        self.assertEqual(secretkey, u'SECRET$6$KEY')
 
     @unittest.skip("")
     def test_4_save_crypto(self):