From 1d9cf2fcaafecdff518a0ef016d793f1641fdc91 Mon Sep 17 00:00:00 2001
From: Abhijit Roy <dalvath3700@gmail.com>
Date: Wed, 5 Jun 2024 15:31:21 +0530
Subject: [PATCH] Incorporate save function from auto-utils into Load key
 function & tests

---
 packages/auto-id/src/helper.ts               |  8 ++++++++
 packages/auto-id/src/keyManagement.ts        | 20 ++++++++++++--------
 packages/auto-id/tests/keyManagement.test.ts |  8 ++++----
 3 files changed, 24 insertions(+), 12 deletions(-)
 create mode 100644 packages/auto-id/src/helper.ts

diff --git a/packages/auto-id/src/helper.ts b/packages/auto-id/src/helper.ts
new file mode 100644
index 00000000..73df1873
--- /dev/null
+++ b/packages/auto-id/src/helper.ts
@@ -0,0 +1,8 @@
+// Helper function to safely parse JSON with a fallback
+export function tryParseJson(jsonString: string, fallback: string): string {
+  try {
+    return JSON.parse(jsonString);
+  } catch (e) {
+    return fallback; // Return the original string if it's not JSON
+  }
+}
\ No newline at end of file
diff --git a/packages/auto-id/src/keyManagement.ts b/packages/auto-id/src/keyManagement.ts
index 7b7fa5b3..8e394d9c 100644
--- a/packages/auto-id/src/keyManagement.ts
+++ b/packages/auto-id/src/keyManagement.ts
@@ -1,7 +1,7 @@
 import { save } from '@autonomys/auto-utils'
 import { KeyObject, createPrivateKey, createPublicKey, generateKeyPairSync } from 'crypto'
 import { promises as fs } from 'fs'
-import { writeFile } from 'fs/promises'
+import { tryParseJson } from './helper'
 
 /**
  * Generates an RSA key pair.
@@ -113,7 +113,7 @@ export function keyToPem(key: KeyObject, password?: string): string {
 export async function saveKey(key: KeyObject, filePath: string, password?: string): Promise<void> {
   try {
     const pem = keyToPem(key, password)
-    await save(filePath, pem);
+    await save(filePath, pem)
   } catch (e: any) {
     throw new Error(`Failed to save key: ${e.message}`)
   }
@@ -167,9 +167,11 @@ export function pemToPrivateKey(pemData: string, password?: string): KeyObject {
  */
 export async function loadPrivateKey(filePath: string, password?: string): Promise<KeyObject> {
   try {
-    const keyData = await fs.readFile(filePath)
-    const privateKey = pemToPrivateKey(keyData.toString(), password)
-    return privateKey
+    let keyData = await fs.readFile(filePath, {encoding: 'utf-8'})
+    // Check if keyData is JSON-encoded and parse it
+    keyData = tryParseJson(keyData, keyData) // Fallback to original data if not JSON
+    const privateKey = pemToPrivateKey(keyData, password)
+    return privateKey;
   } catch (error: any) {
     throw new Error(`Failed to load private key: ${error.message}`)
   }
@@ -213,8 +215,10 @@ export function pemToPublicKey(pemData: string): KeyObject {
  */
 export async function loadPublicKey(filePath: string): Promise<KeyObject> {
   try {
-    const keyData = await fs.readFile(filePath)
-    const publicKey = pemToPublicKey(keyData.toString())
+    let keyData = await fs.readFile(filePath, { encoding: 'utf8' })
+    // Check if keyData is JSON-encoded and parse it
+    keyData = tryParseJson(keyData, keyData); // Fallback to original data if not JSON
+    const publicKey = pemToPublicKey(keyData);
     return publicKey
   } catch (error: any) {
     throw new Error(`Failed to load public key: ${error.message}`)
@@ -288,4 +292,4 @@ export function doPublicKeysMatch(publicKey1: KeyObject, publicKey2: KeyObject):
 
   // Compare the serialized public key data
   return publicKey1Der.equals(publicKey2Der)
-}
+}
\ No newline at end of file
diff --git a/packages/auto-id/tests/keyManagement.test.ts b/packages/auto-id/tests/keyManagement.test.ts
index 72fd2bc2..c343c0c3 100644
--- a/packages/auto-id/tests/keyManagement.test.ts
+++ b/packages/auto-id/tests/keyManagement.test.ts
@@ -139,7 +139,7 @@ describe('PEM to Private/Public key for', () => {
   }
 })
 
-describe('saveKey function', () => {
+describe('Save Key', () => {
   const keyGenerators = [
     { name: 'RSA', generator: generateRsaKeyPair },
     { name: 'Ed25519', generator: generateEd25519KeyPair },
@@ -174,7 +174,7 @@ describe('saveKey function', () => {
         const fileContents = await fs.readFile(filePath, { encoding: 'utf8' })
 
         // Check if the PEM string matches expected, considering JSON.stringify use
-        expect(fileContents).toBe(JSON.stringify(keyToPem(privateKeyObject)));
+        expect(fileContents).toBe(JSON.stringify(keyToPem(privateKeyObject)))
       })
 
       test('should save an encrypted private key to a file', async () => {
@@ -185,7 +185,7 @@ describe('saveKey function', () => {
         const fileContents = await fs.readFile(filePath, { encoding: 'utf8' })
 
         // Parse it back to normal string
-        const actualPemContent = JSON.parse(fileContents);
+        const actualPemContent = JSON.parse(fileContents)
 
         // Check if the file content starts and ends with the expected encrypted private key headers
         expect(actualPemContent.startsWith('-----BEGIN ENCRYPTED PRIVATE KEY-----')).toBe(true)
@@ -200,7 +200,7 @@ describe('saveKey function', () => {
   }
 })
 
-describe('Key loading functions', () => {
+describe('Load Key', () => {
   const keyGenerators = [
     { name: 'RSA', generator: generateRsaKeyPair },
     { name: 'Ed25519', generator: generateEd25519KeyPair },