Merge pull request #61 from Satti-Gowtham/chore/seperate_private_key

Storing Private Key in a seperate file instead of .env File

.gitignore

@@ -161,4 +161,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-.DS_Store
+.DS_Store
+
+*.pem

naptha_sdk/client/hub.py

@@ -1,4 +1,8 @@
 import os
+from naptha_sdk.utils import add_credentials_to_env, get_logger, write_private_key_to_file
+from naptha_sdk.user import generate_keypair
+from naptha_sdk.user import get_public_key, is_hex
+from surrealdb import Surreal
 import traceback
 from typing import Dict, List, Optional, Tuple
 
@@ -87,15 +91,18 @@ async def get_user_by_username(self, username: str) -> Optional[Dict]:
             "SELECT * FROM user WHERE username = $username LIMIT 1",
             {"username": username}
         )
+
         if result and result[0]["result"]:
             return result[0]["result"][0]
+
         return None
 
     async def get_user_by_public_key(self, public_key: str) -> Optional[Dict]:
         result = await self.surrealdb.query(
             "SELECT * FROM user WHERE public_key = $public_key LIMIT 1",
             {"public_key": public_key}
         )
+
         if result and result[0]["result"]:
             return result[0]["result"][0]
         return None
@@ -276,21 +283,28 @@ async def user_setup_flow(hub_url, public_key):
                         print(f"Username '{username}' already exists. Please choose a different username.")
                         continue
                     password = input("Enter password: ")
-                    public_key, private_key = generate_keypair()
+                    public_key, private_key_path = generate_keypair(f"{username}.pem")
                     print(f"Signing up user: {username} with public key: {public_key}")
                     success, token, user_id = await hub.signup(username, password, public_key)
                     if success:
-                        add_credentials_to_env(username, password, private_key)
+                        add_credentials_to_env(username, password, private_key_path)
                         logger.info("Sign up successful!")
                         return token, user_id
                     else:
                         logger.error("Sign up failed. Please try again.")
 
             case None, True, True, False:
-                # User doesn't exist but credentials are provided
-                print(f"Using user credentials in .env. Signing up user: {username} with public key: {public_key}")
+                logger.info("User doesn't exist but some credentials are provided in .env. Using them to create new user.")
+                private_key_path = None
+                if not public_key:
+                    logger.info("No public key provided. Generating new keypair...")
+                    public_key, private_key_path = generate_keypair(f"{username}.pem")
+
+                print(f"Signing up user: {username} with public key: {public_key}")
                 success, token, user_id = await hub.signup(username, password, public_key)
                 if success:
+                    if private_key_path:
+                        add_credentials_to_env(username, password, private_key_path)
                     logger.info("Sign up successful!")
                     return token, user_id
                 else:
@@ -300,6 +314,9 @@ async def user_setup_flow(hub_url, public_key):
             case dict(), True, True, False:
                 # User exists, attempt to sign in
                 logger.info("Using user credentials in .env. User exists. Attempting to sign in...")
+                if os.getenv("PRIVATE_KEY") and is_hex(os.getenv("PRIVATE_KEY")):
+                    write_private_key_to_file(os.getenv("PRIVATE_KEY"), username)
+
                 success, token, user_id = await hub.signin(username, password)
                 if success:
                     logger.info("Sign in successful!")
@@ -310,5 +327,4 @@ async def user_setup_flow(hub_url, public_key):
 
             case _:
                 logger.error("Unexpected case encountered in user setup flow.")
-                raise Exception("Unexpected error in user setup. Please check your configuration and try again.")
-
+                raise Exception("Unexpected error in user setup. Please check your configuration and try again.")

naptha_sdk/user.py

@@ -1,11 +1,40 @@
 from ecdsa import SigningKey, SECP256k1
+import os, re
 
-def generate_keypair(private_key=None): 
+def generate_keypair(private_key_filename=None): 
     private_key = SigningKey.generate(curve=SECP256k1).to_string().hex()
-    public_key = get_public_key(private_key)
-    return public_key, private_key
+    public_key = generate_public_key(private_key)
+    pkfile = private_key_filename if private_key_filename else 'private_key.pem'
 
-def get_public_key(private_key_hex):
+    # Save the private key to a file
+    with open(os.path.join(os.getcwd(), pkfile), 'w') as file:
+        file.write(private_key)
+
+    return public_key, os.path.join(os.getcwd(), pkfile)
+
+def get_public_key(private_key):
+    # To ensure old users can still login
+    if private_key and is_hex(private_key):
+        private_key_hex = private_key
+    elif private_key and os.path.isfile(private_key):
+        with open(private_key) as file:
+            content = file.read().strip()
+
+            if content:
+                private_key_hex = content
+            else:
+                return None
+    else:
+        return None
+
+    return generate_public_key(private_key_hex)
+
+def generate_public_key(private_key_hex):
     private_key = SigningKey.from_string(bytes.fromhex(private_key_hex), curve=SECP256k1)
     public_key = private_key.get_verifying_key()
-    return public_key.to_string().hex()
+
+    return public_key.to_string().hex()
+
+def is_hex(string):
+    # Check if the string matches the pattern for a hexadecimal key
+    return bool(re.match(r'^[0-9a-fA-F]{64}$', string))

naptha_sdk/utils.py

@@ -21,7 +21,7 @@ def load_yaml(cfg_path):
         cfg = yaml.load(file, Loader=yaml.FullLoader)
     return cfg
 
-def add_credentials_to_env(username, password, private_key):
+def add_credentials_to_env(username, password, private_key_path):
     env_file_path = os.path.join(os.getcwd(), '.env')
     updated_lines = []
     hub_user_found = False
@@ -38,7 +38,7 @@ def add_credentials_to_env(username, password, private_key):
                 updated_lines.append(f"HUB_PASSWORD={password}\n")
                 hub_pass_found = True
             elif line.startswith('PRIVATE_KEY='):
-                updated_lines.append(f"PRIVATE_KEY={private_key}\n")
+                updated_lines.append(f"PRIVATE_KEY={private_key_path}\n")
                 private_key_found = True
             else:
                 updated_lines.append(line)
@@ -49,13 +49,41 @@ def add_credentials_to_env(username, password, private_key):
     if not hub_pass_found:
         updated_lines.append(f"HUB_PASSWORD={password}\n")
     if not private_key_found:
-        updated_lines.append(f"PRIVATE_KEY={private_key}\n")
+        updated_lines.append(f"PRIVATE_KEY={private_key_path}\n")
 
     # Write the updated content back to the .env file
     with open(env_file_path, 'w') as env_file:
         env_file.writelines(updated_lines)
 
-    print("Your credentials and private key have been updated in the .env file. You can now use these credentials to authenticate in future sessions.")
+    print("Your credentials have been updated in the .env file. You can now use these credentials to authenticate in future sessions.")
+
+def write_private_key_to_file(private_key, username):
+    private_key_file_path = os.path.join(os.getcwd(), f'{username}.pem')
+    with open(private_key_file_path, 'w') as file:
+        file.write(private_key)
+
+    update_private_key_in_env(private_key_file_path)    
+
+def update_private_key_in_env(private_key_path):
+    env_file_path = os.path.join(os.getcwd(), '.env')
+    updated_lines = []
+    private_key_found = False
+
+    with open(env_file_path, 'r') as env_file:
+        for line in env_file:
+            if line.startswith('PRIVATE_KEY='):
+                updated_lines.append(f"PRIVATE_KEY={private_key_path}\n")
+                private_key_found = True
+            else:
+                updated_lines.append(line)
+
+    if not private_key_found:
+        updated_lines.append(f"PRIVATE_KEY={private_key_path}\n")
+
+    with open(env_file_path, 'w') as env_file:
+        env_file.writelines(updated_lines)
+
+    print("Your private key have been updated in the .env file")
 
 class AsyncMixin:
     def __init__(self, *args, **kwargs):