From c82776aed6006c253b9cd3c7000b2c39e33baa7b Mon Sep 17 00:00:00 2001
From: Jeremy Frank <jeremy@subspace.network>
Date: Fri, 5 Apr 2024 08:21:22 -0600
Subject: [PATCH] add helper functions to key_management

---
 auto_identity/__init__.py       | 17 ++++++++-
 auto_identity/key_management.py | 67 +++++++++++++++++++++++++--------
 setup.py                        |  2 +-
 3 files changed, 68 insertions(+), 18 deletions(-)

diff --git a/auto_identity/__init__.py b/auto_identity/__init__.py
index c616cd5..83c1a96 100644
--- a/auto_identity/__init__.py
+++ b/auto_identity/__init__.py
@@ -6,12 +6,21 @@
 """
 
 from substrateinterface import Keypair
-from .key_management import generate_rsa_key_pair, generate_ed25519_key_pair, key_to_hex, load_private_key, load_public_key
+from .key_management import (
+    generate_rsa_key_pair,
+    generate_ed25519_key_pair,
+    key_to_hex,
+    key_to_pem,
+    pem_to_private_key,
+    load_private_key,
+    pem_to_public_key,
+    load_public_key,
+    save_key)
 from .certificate_management import create_csr, issue_certificate, self_issue_certificate, get_subject_common_name
 from .registry import Registry
 from .utils import der_encode_signature_algorithm_oid
 
-__version__ = '0.1.1'
+__version__ = '0.1.2'
 
 __all__ = [
     "generate_rsa_key_pair",
@@ -19,6 +28,10 @@
     "key_to_hex",
     "load_private_key",
     "load_public_key",
+    "save_key",
+    "pem_to_private_key",
+    "pem_to_public_key",
+    "key_to_pem",
     "create_csr",
     "issue_certificate",
     "self_issue_certificate",
diff --git a/auto_identity/key_management.py b/auto_identity/key_management.py
index 7c65b42..f9ddbbe 100644
--- a/auto_identity/key_management.py
+++ b/auto_identity/key_management.py
@@ -34,15 +34,13 @@ def generate_ed25519_key_pair() -> tuple:
     return private_key, public_key
 
 
-def save_key(key, file_path: str, password: str = None) -> None:
+def key_to_pem(key, password: str = None) -> str:
     """
-    Saves a private or public key to a file. If it's a private key and a password is provided,
-    the key will be encrypted.
+    Converts a private or public key to a PEM string.
 
     Args:
-        key: The key to save (private or public).
-        file_path (str): Path to the file where the key should be saved.
-        password (str): Optional password to encrypt the private key.
+        key: The key to convert (private or public).
+        password (str): The password used to encrypt the key.
     """
     if hasattr(key, 'private_bytes'):
         encoding = serialization.Encoding.PEM
@@ -55,10 +53,42 @@ def save_key(key, file_path: str, password: str = None) -> None:
         format = serialization.PublicFormat.SubjectPublicKeyInfo
         key_data = key.public_bytes(encoding, format)
 
+    return key_data
+
+
+def save_key(key, file_path: str, password: str = None) -> None:
+    """
+    Saves a private or public key to a file. If it's a private key and a password is provided,
+    the key will be encrypted.
+
+    Args:
+        key: The key to save (private or public).
+        file_path (str): Path to the file where the key should be saved.
+        password (str): Optional password to encrypt the private key.
+    """
+    key_data = key_to_pem(key, password)
+
     with open(file_path, "wb") as key_file:
         key_file.write(key_data)
 
 
+def pem_to_private_key(pem_data: str, password: str = None):
+    """
+    Converts a PEM string to a private or public key. If the PEM string is encrypted, a password must be provided.
+
+    Args:
+        pem_data (str): The PEM string to convert.
+        password (str): The password used to encrypt the key.
+    """
+
+    private_key = serialization.load_pem_private_key(
+        pem_data,
+        password=password.encode() if password else None,
+        backend=default_backend()
+    )
+    return private_key
+
+
 def load_private_key(file_path: str, password: str = None):
     """
     Loads a private key from a file. If the file is encrypted, a password must be provided.
@@ -71,14 +101,24 @@ def load_private_key(file_path: str, password: str = None):
         The private key.
     """
     with open(file_path, "rb") as key_file:
-        private_key = serialization.load_pem_private_key(
-            key_file.read(),
-            password=password.encode() if password else None,
-            backend=default_backend()
-        )
+        private_key = pem_to_private_key(key_file.read(), password)
     return private_key
 
 
+def pem_to_public_key(pem_data: str):
+    """
+    Converts a PEM string to a public key.
+
+    Args:
+        pem_data (str): The PEM string to convert.
+    """
+    public_key = serialization.load_pem_public_key(
+        pem_data,
+        backend=default_backend()
+    )
+    return public_key
+
+
 def load_public_key(file_path: str):
     """
     Loads a public key from a file.
@@ -90,10 +130,7 @@ def load_public_key(file_path: str):
         The public key.
     """
     with open(file_path, "rb") as key_file:
-        public_key = serialization.load_pem_public_key(
-            key_file.read(),
-            backend=default_backend()
-        )
+        public_key = pem_to_public_key(key_file.read())
     return public_key
 
 
diff --git a/setup.py b/setup.py
index 5f60571..3373f26 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='auto-sdk',
-    version='0.1.1',
+    version='0.1.2',
     author='Autonomys',
     author_email='jeremy@subspace.network',
     url='https://github.com/subspace/auto-kit',