clean up some of the enc key stuff

demo.py

@@ -17,9 +17,8 @@ import apns
 import ids
 from ids.keydec import IdentityKeys
 
-FORMAT = "%(message)s"
 logging.basicConfig(
-    level="NOTSET", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()]
+    level=logging.NOTSET, format="%(message)s", datefmt="[%X]", handlers=[RichHandler()]
 )
 
 # Set sane log levels
@@ -67,28 +66,22 @@ else:
 
     user.authenticate(username, password)
 
-# Generate a new RSA keypair for the identity
-# Load the old keypair if it exists
-if CONFIG.get("encrypt") is not None:
-    priv_enc = load_pem_private_key(CONFIG["encrypt"].encode(), password=None)
-else:
-    priv_enc = rsa.generate_private_key(public_exponent=65537, key_size=1280)
-pub_enc = priv_enc.public_key()
+user.ec_key = CONFIG.get("encryption", {}).get("ec_key")
+user.rsa_key = CONFIG.get("encryption", {}).get("rsa_key")
 
-if CONFIG.get("id", {}).get("cert") is not None:
+if CONFIG.get("id", {}).get("cert") is not None and user.ec_key is not None and user.rsa_key is not None:
     id_keypair = ids._helpers.KeyPair(CONFIG["id"]["key"], CONFIG["id"]["cert"])
     user.restore_identity(id_keypair)
 else:
-    # vd = input_multiline("Enter validation data: ")
+    logging.info("Registering new identity...")
     import emulated.nac
 
     vd = emulated.nac.generate_validation_data()
     vd = b64encode(vd).decode()
 
-    published_keys = IdentityKeys(None, pub_enc)
-    user.register(vd, published_keys)
+    user.register(vd)
 
-logging.info("Waiting for incomming messages...")
+logging.info("Waiting for incoming messages...")
 
 # Create a thread to send keepalive messages
 
@@ -102,6 +95,10 @@ def keepalive():
 threading.Thread(target=keepalive, daemon=True).start()
 
 # Write config.json
+CONFIG["encryption"] = {
+    "ec_key": user.ec_key,
+    "rsa_key": user.rsa_key,
+}
 CONFIG["id"] = {
     "key": user._id_keypair.key,
     "cert": user._id_keypair.cert,
@@ -118,20 +115,10 @@ CONFIG["push"] = {
     "cert": user.push_connection.cert,
 }
 
-
-CONFIG["encrypt"] = (
-    priv_enc.private_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PrivateFormat.TraditionalOpenSSL,
-        encryption_algorithm=serialization.NoEncryption(),
-    )
-    .decode("utf-8")
-    .strip()
-)
-
 with open("config.json", "w") as f:
     json.dump(CONFIG, f, indent=4)
 
+user_rsa_key = load_pem_private_key(user.rsa_key.encode(), password=None)
 
 def decrypt(payload):
     # print(payload[1:3])
@@ -140,7 +127,7 @@ def decrypt(payload):
     payload = payload[3 : length + 3]
     # print("Decrypting payload", payload)
 
-    decrypted1 = priv_enc.decrypt(
+    decrypted1 = user_rsa_key.decrypt(
         payload[:160],
         padding.OAEP(
             mgf=padding.MGF1(algorithm=hashes.SHA1()),

ids/__init__.py

@@ -28,6 +28,8 @@ class IDSUser:
             self.push_connection.private_key, self.push_connection.cert
         )
 
+        self.ec_key = self.rsa_key = None
+
     def __str__(self):
         return f"IDSUser(user_id={self.user_id}, handles={self.handles}, push_token={b64encode(self.push_connection.token).decode()})"
 
@@ -53,7 +55,16 @@ class IDSUser:
         self.handles = handles 
 
     # This is a separate call so that the user can make sure the first part succeeds before asking for validation data
-    def register(self, validation_data: str, published_keys: keydec.IdentityKeys):
+    def register(self, validation_data: str):
+        """
+        self.ec_key, self.rsa_key will be set to a randomly gnenerated EC and RSA keypair
+        if they are not already set
+        """
+        if self.ec_key is None or self.rsa_key is None:
+            self.ec_key, self.rsa_key, published_keys = keydec.generate_keys()
+        else:
+            published_keys = keydec.load_keys(self.ec_key, self.rsa_key)
+        
         cert = identity.register(
             b64encode(self.push_connection.token),
             self.handles,

ids/keydec.py

@@ -1,10 +1,36 @@
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives import serialization
 
 from base64 import b64decode, b64encode
 
 from io import BytesIO
 
+def generate_keys() -> tuple[str, str, 'IdentityKeys']:
+    """
+    ECDSA key, RSA key, IdentityKeys
+    """
+    ecdsa_key = ec.generate_private_key(ec.SECP256R1())
+    rsa_key = rsa.generate_private_key(65537, 1280)
+
+    # Serialize the keys into PEM
+    ecdsa_key_p = ecdsa_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption(),
+    ).decode("utf-8").strip()
+    rsa_key_p = rsa_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption(),
+    ).decode("utf-8").strip()
+    return ecdsa_key_p, rsa_key_p, IdentityKeys(ecdsa_key.public_key(), rsa_key.public_key())
+
+def load_keys(ecdsa_key_p: str, rsa_key_p: str) -> 'IdentityKeys':
+    ecdsa_key = serialization.load_pem_private_key(ecdsa_key_p.encode(), password=None)
+    rsa_key = serialization.load_pem_private_key(rsa_key_p.encode(), password=None)
+    return IdentityKeys(ecdsa_key.public_key(), rsa_key.public_key())
+
 class IdentityKeys():
     def __init__(self, ecdsa_key: ec.EllipticCurvePublicKey, rsa_key: rsa.RSAPublicKey):
         self.ecdsa_key = ecdsa_key