From 330b3b99cf2effea09dcfeb8951a03d5bdf8cc21 Mon Sep 17 00:00:00 2001
From: emjshrx <emjshrx@gmail.com>
Date: Sat, 3 Feb 2024 16:26:09 +0530
Subject: [PATCH] feat: encrypt and decrypt keys on storage

---
 src/wallet/db/db.interface.ts | 10 ++++++++--
 src/wallet/wallet.ts          | 35 +++++++++++++++++++++++++++++++----
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/src/wallet/db/db.interface.ts b/src/wallet/db/db.interface.ts
index e29140a..08a9828 100644
--- a/src/wallet/db/db.interface.ts
+++ b/src/wallet/db/db.interface.ts
@@ -7,8 +7,14 @@ export type DbInterface = {
     getStatus(): string;
     getVersion(): Promise<number>;
     setVersion(version: number): Promise<void>;
-    getMasterKey(): Promise<{ privateKey: Buffer; chaincode: Buffer }>;
-    setMasterKey(privateKey: Buffer, chaincode: Buffer): Promise<void>;
+    getEncryptedMasterKey(): Promise<{
+        encryptedKey: string;
+        chaincode: Buffer;
+    }>;
+    setEncryptedMasterKey(
+        encryptedKey: string,
+        chaincode: Buffer,
+    ): Promise<void>;
     saveAddress(address: string, path: string): Promise<void>;
     getAddress(address: string): Promise<string>;
     hasAddress(address: string): Promise<boolean>;
diff --git a/src/wallet/wallet.ts b/src/wallet/wallet.ts
index 987d3ba..135c208 100644
--- a/src/wallet/wallet.ts
+++ b/src/wallet/wallet.ts
@@ -11,6 +11,7 @@ import { Coin } from './coin.ts';
 import { ECPairFactory } from 'ecpair';
 import { createOutputs } from '../core';
 import { toXOnly } from 'bitcoinjs-lib/src/psbt/bip371';
+import bip38 from 'bip38';
 
 initEccLib(ecc);
 const ECPair = ECPairFactory(ecc);
@@ -27,25 +28,44 @@ export class Wallet {
     private masterKey: BIP32Interface;
     private receiveDepth: number = 0;
     private changeDepth: number = 0;
+    private encryptionPassword: string = '12345678';
+    private compressedKey: boolean = true;
 
     constructor(config: WalletConfigOptions) {
         this.db = config.db;
         this.network = config.networkClient;
     }
 
-    async init(mnemonic?: string) {
+    async init(mnemonic?: string, password?: string) {
+        if (password) {
+            this.encryptionPassword = password;
+        }
+
         await this.db.open();
 
         if (mnemonic) {
             const seed = mnemonicToSeedSync(mnemonic).toString('hex');
             this.masterKey = bip32.fromSeed(Buffer.from(seed, 'hex'));
-            await this.db.setMasterKey(
+            const encryptedKey = bip38.encrypt(
                 this.masterKey.privateKey,
+                this.compressedKey,
+                this.encryptionPassword,
+            );
+            await this.db.setEncryptedMasterKey(
+                encryptedKey,
                 this.masterKey.chainCode,
             );
         } else {
-            const { privateKey, chaincode } = await this.db.getMasterKey();
-            this.masterKey = bip32.fromPrivateKey(privateKey, chaincode);
+            const { encryptedKey, chaincode } =
+                await this.db.getEncryptedMasterKey();
+            const decryptedKey = bip38.decrypt(
+                encryptedKey,
+                this.encryptionPassword,
+            );
+            this.masterKey = bip32.fromPrivateKey(
+                decryptedKey.privateKey,
+                chaincode,
+            );
         }
     }
 
@@ -258,4 +278,11 @@ export class Wallet {
 
         return tx.getId();
     }
+    setPassword(newPassword: string, oldPassword: string): void {
+        if (oldPassword !== this.encryptionPassword) {
+            throw new Error('Wrong old encryption password!');
+        } else {
+            this.encryptionPassword = newPassword;
+        }
+    }
 }