feat: encrypt and decrypt keys on storage

Bitshala-Incubator · Feb 9, 2024 · ae2467e · ae2467e
1 parent de709ca
commit ae2467e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 21 deletions.
diff --git a/src/wallet/db/db.interface.ts b/src/wallet/db/db.interface.ts
@@ -1,4 +1,3 @@
-import { Buffer } from 'buffer';
 import { Coin } from '../coin.ts';
 
 export type DbInterface = {
@@ -7,8 +6,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>;
+    getMasterKey(): Promise<{
+        encryptedPrivateKey: string;
+        encryptedChainCode: string;
+    }>;
+    setMasterKey(
+        encryptedPrivateKey: string,
+        encryptedChainCode: string,
+    ): Promise<void>;
     saveAddress(address: string, path: string): Promise<void>;
     getAddress(address: string): Promise<string>;
     hasAddress(address: string): Promise<boolean>;

diff --git a/src/wallet/db/level/db.ts b/src/wallet/db/level/db.ts
@@ -1,7 +1,6 @@
 import { Level } from 'level';
 import { wdb } from './layout.ts';
 import { DbInterface } from '../db.interface.ts';
-import { Buffer } from 'buffer';
 import { Coin } from '../../coin.ts';
 
 export type LevelDBConfigOptions = {
@@ -39,24 +38,20 @@ export class WalletDB implements DbInterface {
         await this.db.put(wdb.V, version.toString());
     }
 
-    public async getMasterKey(): Promise<{
-        privateKey: Buffer;
-        chaincode: Buffer;
-    }> {
+    public async getMasterKey() {
         const masterKey = await this.db.get(wdb.M);
-        const privateKey = Buffer.from(masterKey.slice(0, 64), 'hex');
-        const chaincode = Buffer.from(masterKey.slice(64), 'hex');
+        const [encryptedPrivateKey, encryptedChainCode] = masterKey.split(':');
 
-        return { privateKey, chaincode };
+        return { encryptedPrivateKey, encryptedChainCode };
     }
 
     public async setMasterKey(
-        privateKey: Buffer,
-        chaincode: Buffer,
+        encryptedPrivateKey: string,
+        encryptedChainCode: string,
     ): Promise<void> {
         await this.db.put(
             wdb.M,
-            Buffer.concat([privateKey, chaincode]).toString('hex'),
+            `${encryptedPrivateKey}:${encryptedChainCode}`,
         );
     }
 

diff --git 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, encodeSilentPaymentAddress } from '../core';
 import { toXOnly } from 'bitcoinjs-lib/src/psbt/bip371';
+import { encrypt, decrypt } from 'bip38';
 
 initEccLib(ecc);
 const ECPair = ECPairFactory(ecc);
@@ -21,6 +22,8 @@ export type WalletConfigOptions = {
     networkClient: NetworkInterface;
 };
 
+const DEFAULT_ENCRYPTION_PASSWORD = '12345678';
+
 export class Wallet {
     private readonly db: DbInterface;
     private readonly network: NetworkInterface;
@@ -33,19 +36,29 @@ export class Wallet {
         this.network = config.networkClient;
     }
 
-    async init(mnemonic?: string) {
+    async init(params?: { mnemonic?: string; password?: string }) {
+        const { mnemonic, password } = params;
         await this.db.open();
 
         if (mnemonic) {
             const seed = mnemonicToSeedSync(mnemonic).toString('hex');
             this.masterKey = bip32.fromSeed(Buffer.from(seed, 'hex'));
-            await this.db.setMasterKey(
-                this.masterKey.privateKey,
-                this.masterKey.chainCode,
-            );
+            this.writeMasterKey(password ?? DEFAULT_ENCRYPTION_PASSWORD);
         } else {
-            const { privateKey, chaincode } = await this.db.getMasterKey();
-            this.masterKey = bip32.fromPrivateKey(privateKey, chaincode);
+            const { encryptedPrivateKey, encryptedChainCode } =
+                await this.db.getMasterKey();
+            const { privateKey: decryptedPrivateKey } = decrypt(
+                encryptedPrivateKey,
+                password ?? DEFAULT_ENCRYPTION_PASSWORD,
+            );
+            const { privateKey: decryptedChainCode } = decrypt(
+                encryptedChainCode,
+                password ?? DEFAULT_ENCRYPTION_PASSWORD,
+            );
+            this.masterKey = bip32.fromPrivateKey(
+                decryptedPrivateKey,
+                decryptedChainCode,
+            );
         }
     }
 
@@ -56,6 +69,26 @@ export class Wallet {
         await this.db.close();
     }
 
+    async writeMasterKey(newPassword: string) {
+        if (!this.masterKey) {
+            throw new Error(
+                'Wallet not initialized. Please call wallet.init()',
+            );
+        } else {
+            const encryptedPrivateKey = encrypt(
+                this.masterKey.privateKey,
+                false,
+                newPassword,
+            );
+            const encryptedChainCode = encrypt(
+                this.masterKey.chainCode,
+                false,
+                newPassword,
+            );
+            await this.db.setMasterKey(encryptedPrivateKey, encryptedChainCode);
+        }
+    }
+
     private async deriveAddress(path: string): Promise<string> {
         const child = this.masterKey.derivePath(path);
         const { address } = payments.p2wpkh({