From 452f4285eb7d7e374710cbe0a3401520042b7c63 Mon Sep 17 00:00:00 2001
From: Franck Royer <franck@status.im>
Date: Thu, 24 Jun 2021 16:05:22 +1000
Subject: [PATCH] Move most EthCrypto usage to crypto.ts

---
 examples/eth-dm/src/App.tsx         | 16 ++++---------
 examples/eth-dm/src/SendMessage.tsx |  4 ++--
 examples/eth-dm/src/crypto.ts       | 35 ++++++++++++++++++++++++++---
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/examples/eth-dm/src/App.tsx b/examples/eth-dm/src/App.tsx
index b94b5f9397..8add3f6467 100644
--- a/examples/eth-dm/src/App.tsx
+++ b/examples/eth-dm/src/App.tsx
@@ -7,11 +7,12 @@ import { ethers } from 'ethers';
 import { Web3Provider } from '@ethersproject/providers';
 import {
   createPublicKeyMessage,
+  decryptMessage,
   generateEthDmKeyPair,
   KeyPair,
+  recoverKeysFromPrivateKey,
   validatePublicKeyMessage,
 } from './crypto';
-import * as EthCrypto from 'eth-crypto';
 import { decode, DirectMessage, encode, PublicKeyMessage } from './messages';
 import { Message, Messages } from './Messages';
 import 'fontsource-roboto';
@@ -225,10 +226,7 @@ async function handleDirectMessage(
   console.log('Waku Message received:', wakuMsg);
   if (!wakuMsg.payload) return;
   const directMessage: DirectMessage = decode(wakuMsg.payload);
-  const text = await EthCrypto.decryptWithPrivateKey(
-    privateKey,
-    directMessage.encMessage
-  );
+  const text = await decryptMessage(privateKey, directMessage);
 
   const timestamp = wakuMsg.timestamp ? wakuMsg.timestamp : new Date();
 
@@ -251,13 +249,7 @@ function saveKeysToStorage(ethDmKeyPair: KeyPair) {
 function retrieveKeysFromStorage() {
   const privateKey = window.localStorage.getItem(EthDmKeyStorageKey);
   if (privateKey) {
-    const publicKey = EthCrypto.publicKeyByPrivateKey(privateKey);
-    const address = EthCrypto.publicKey.toAddress(publicKey);
-    return {
-      privateKey,
-      publicKey,
-      address,
-    };
+    return recoverKeysFromPrivateKey(privateKey);
   }
   return;
 }
diff --git a/examples/eth-dm/src/SendMessage.tsx b/examples/eth-dm/src/SendMessage.tsx
index f35bd2c4de..28504a2fbb 100644
--- a/examples/eth-dm/src/SendMessage.tsx
+++ b/examples/eth-dm/src/SendMessage.tsx
@@ -8,9 +8,9 @@ import {
 } from '@material-ui/core';
 import React, { ChangeEvent, useState, KeyboardEvent } from 'react';
 import { Waku, WakuMessage } from 'js-waku';
-import * as EthCrypto from 'eth-crypto';
 import { DirectMessage, encode } from './messages';
 import { DirectMessageContentTopic } from './App';
+import { encryptMessage } from './crypto';
 
 const useStyles = makeStyles((theme) => ({
   formControl: {
@@ -111,7 +111,7 @@ async function encodeEncryptedWakuMessage(
   publicKey: string,
   address: string
 ): Promise<WakuMessage> {
-  const encryptedMsg = await EthCrypto.encryptWithPublicKey(publicKey, message);
+  const encryptedMsg = await encryptMessage(publicKey, message);
 
   const directMsg: DirectMessage = {
     toAddress: address,
diff --git a/examples/eth-dm/src/crypto.ts b/examples/eth-dm/src/crypto.ts
index 4ed906d169..c4d87e8ee4 100644
--- a/examples/eth-dm/src/crypto.ts
+++ b/examples/eth-dm/src/crypto.ts
@@ -3,7 +3,7 @@ import '@ethersproject/shims';
 import * as EthCrypto from 'eth-crypto';
 import { ethers } from 'ethers';
 import { Signer } from '@ethersproject/abstract-signer';
-import { PublicKeyMessage } from './messages';
+import { DirectMessage, PublicKeyMessage } from './messages';
 
 export interface KeyPair {
   privateKey: string;
@@ -59,8 +59,37 @@ export function validatePublicKeyMessage(msg: PublicKeyMessage): boolean {
  * context.
  */
 function formatPublicKeyForSignature(ethDmPublicKey: string): string {
-  const txt = JSON.stringify({
+  return JSON.stringify({
     ethDmPublicKey,
   });
-  return txt;
+}
+
+/**
+ * Decrypt a Direct Message using the private key.
+ */
+export function decryptMessage(
+  privateKey: string,
+  directMessage: DirectMessage
+) {
+  return EthCrypto.decryptWithPrivateKey(privateKey, directMessage.encMessage);
+}
+
+/**
+ * Recover Public Key and address from Private Key
+ */
+export function recoverKeysFromPrivateKey(privateKey: string) {
+  const publicKey = EthCrypto.publicKeyByPrivateKey(privateKey);
+  const address = EthCrypto.publicKey.toAddress(publicKey);
+  return {
+    privateKey,
+    publicKey,
+    address,
+  };
+}
+
+/**
+ * Encrypt message with given Public Key
+ */
+export async function encryptMessage(publicKey: string, message: string) {
+  return await EthCrypto.encryptWithPublicKey(publicKey, message);
 }