Merge pull request #117 from tuxxy/frag-keystore

Add KeyFrag to KeyStore

nkms/keystore/constants.py

@@ -1,5 +1,8 @@
 ENC_KEYPAIR_BYTE = b'E'
 SIG_KEYPAIR_BYTE = b'S'
 
+REKEY_FRAG_ID_LEN = 32
+REKEY_FRAG_KEY_LEN = 32
+
 PUB_KEY_BYTE = b'0'
 PRIV_KEY_BYTE = b'1'

nkms/keystore/db/models.py

@@ -24,3 +24,15 @@ class Key(Base):
         :return: Fingerprint of key as a string
         """
         return sha3.keccak_256(self.key_data[2:]).hexdigest().encode()
+
+
+class KeyFrag(Base):
+    __tablename__ = 'keyfrags'
+
+    id = Column(Integer, primary_key=True)
+    hrac = Column(LargeBinary, unique=True)
+    key_frag = Column(LargeBinary, unique=True)
+
+    def __init__(self, hrac, key_frag):
+        self.hrac = hrac
+        self.key_frag = key_frag

nkms/keystore/keystore.py

@@ -1,8 +1,11 @@
 import sha3
 from nkms.keystore import keypairs, constants
-from nkms.keystore.db.models import Key
+from nkms.keystore.db.models import Key, KeyFrag
+from nkms.crypto.utils import BytestringSplitter
+from nkms.crypto.signature import Signature
 from sqlalchemy.orm import sessionmaker
 from typing import Union
+from npre.umbral import RekeyFrag
 
 
 class KeyNotFound(KeyError):
@@ -17,6 +20,12 @@ class KeyStore(object):
     A storage class of cryptographic keys.
     """
 
+    kFrag_splitter = BytestringSplitter(
+            Signature,
+            (bytes, constants.REKEY_FRAG_ID_LEN),
+            (bytes, constants.REKEY_FRAG_KEY_LEN)
+    )
+
     def __init__(self, sqlalchemy_engine=None):
         """
         Initalizes a KeyStore object.
@@ -81,6 +90,28 @@ class KeyStore(object):
                     "No key with fingerprint {} found.".format(fingerprint))
         return keypairs.Keypair.deserialize_key(key.key_data)
 
+    def get_kfrag(self, hrac: bytes, get_sig: bool=False) -> RekeyFrag:
+        """
+        Returns a RekeyFrag from the KeyStore.
+
+        :param hrac: HRAC in bytes
+
+        :return: Deserialized RekeyFrag from KeyStore
+        """
+        kfrag = self.session.query(KeyFrag).filter_by(hrac=hrac).first()
+        if not kfrag:
+            raise KeyNotFound(
+                "No KeyFrag with HRAC {} found."
+                .format(hrac)
+            )
+        # TODO: Make this use a class
+        sig, id, key = self.kFrag_splitter(kfrag.key_frag)
+
+        kFrag = RekeyFrag(id=id, key=key)
+        if get_sig:
+            return (kFrag, sig)
+        return kFrag
+
     def add_key(self,
                 keypair: Union[keypairs.EncryptingKeypair,
                                keypairs.SigningKeypair],
@@ -89,7 +120,6 @@ class KeyStore(object):
         Gets a fingerprint of the key and adds it to the keystore.
 
         :param key: Key, in bytes, to add to lmdb
-        ::
 
         :return: Fingerprint, in bytes, of the added key
         """
@@ -105,6 +135,23 @@ class KeyStore(object):
         self.session.commit()
         return fingerprint
 
+    def add_kfrag(self, hrac: bytes, kfrag: RekeyFrag, sig: bytes=None):
+        """
+        Adds a RekeyFrag to sqlite.
+
+        :param hrac: Hashed Resource Authenticate Code
+        :param kfrag: RekeyFrag instance to add to sqlite
+        :param sig: Signature of kfrag (if exists)
+        """
+        kfrag_data = b''
+        if sig:
+            kfrag_data += sig
+        kfrag_data += kfrag.id + kfrag.key
+
+        kfrag = KeyFrag(hrac, kfrag_data)
+        self.session.add(kfrag)
+        self.session.commit()
+
     def del_key(self, fingerprint: bytes):
         """
         Deletes a key from the KeyStore.
@@ -113,3 +160,12 @@ class KeyStore(object):
         """
         self.session.query(Key).filter_by(fingerprint=fingerprint).delete()
         self.session.commit()
+
+    def del_kfrag(self, hrac: bytes):
+        """
+        Deletes a RekeyFrag from sqlite.
+
+        :param hrac: Hashed Resource Authentication Code
+        """
+        self.session.query(KeyFrag).filter_by(hrac=hrac).delete()
+        self.session.commit()

tests/keystore/test_keystore.py

@@ -3,6 +3,8 @@ import sha3
 from sqlalchemy import create_engine
 from nkms.keystore.db import Base
 from nkms.keystore import keystore, keypairs
+from npre.umbral import RekeyFrag
+from nkms.crypto import api as API
 
 
 class TestKeyStore(unittest.TestCase):
@@ -24,7 +26,7 @@ class TestKeyStore(unittest.TestCase):
         self.assertEqual(bytes, type(keypair.privkey))
         self.assertEqual(bytes, type(keypair.pubkey))
 
-    def test_sqlite_keystore(self):
+    def test_key_sqlite_keystore(self):
         keypair = self.ks.generate_encrypting_keypair()
         self.assertEqual(keypairs.EncryptingKeypair, type(keypair))
         self.assertEqual(bytes, type(keypair.privkey))
@@ -69,3 +71,23 @@ class TestKeyStore(unittest.TestCase):
         self.ks.del_key(fingerprint_priv)
         with self.assertRaises(keystore.KeyNotFound):
             key = self.ks.get_key(fingerprint_priv)
+
+    def test_keyfrag_sqlite(self):
+        rand_sig = API.secure_random(65)
+        rand_id = API.secure_random(32)
+        rand_key = API.secure_random(32)
+        rand_hrac = API.secure_random(32)
+
+        kfrag = RekeyFrag(id=rand_id, key=rand_key)
+        self.ks.add_kfrag(rand_hrac, kfrag, sig=rand_sig)
+
+        # Check that kfrag was added
+        kfrag, signature = self.ks.get_kfrag(rand_hrac, get_sig=True)
+        self.assertEqual(rand_sig, signature)
+        self.assertEqual(kfrag.id, rand_id)
+        self.assertEqual(kfrag.key, rand_key)
+
+        # Check that kfrag gets deleted
+        self.ks.del_kfrag(rand_hrac)
+        with self.assertRaises(keystore.KeyNotFound):
+            key = self.ks.get_key(rand_hrac)