From fd09cffd9ff1854c2b972d23e82ae8037bf5aa72 Mon Sep 17 00:00:00 2001
From: matstyler <mat.franczuk@gmail.com>
Date: Wed, 28 Feb 2024 15:21:08 +0100
Subject: [PATCH 1/2] :sparkles: feat: Encrypt/decrypt using MetaMask

---
 wallets/metamask/src/metamask.ts              | 32 ++++++++++++-----
 .../NotificationPage/actions/encryption.ts    | 10 ++++++
 .../pages/NotificationPage/actions/index.ts   |  1 +
 .../src/pages/NotificationPage/page.ts        | 16 ++++++++-
 .../test/e2e/metamask/encrypt.spec.ts         | 34 +++++++++++++++++++
 5 files changed, 84 insertions(+), 9 deletions(-)
 create mode 100644 wallets/metamask/src/pages/NotificationPage/actions/encryption.ts
 create mode 100644 wallets/metamask/test/e2e/metamask/encrypt.spec.ts

diff --git a/wallets/metamask/src/metamask.ts b/wallets/metamask/src/metamask.ts
index 7268c0304..9e864f1bf 100644
--- a/wallets/metamask/src/metamask.ts
+++ b/wallets/metamask/src/metamask.ts
@@ -384,6 +384,30 @@ export class MetaMask {
     await this.settingsPage.disableEthSign()
   }
 
+  async addNewToken() {
+    if (!this.extensionId) {
+      throw NO_EXTENSION_ID_ERROR
+    }
+
+    await this.notificationPage.addNewToken(this.extensionId)
+  }
+
+  async providePublicEncryptionKey() {
+    if (!this.extensionId) {
+      throw NO_EXTENSION_ID_ERROR
+    }
+
+    await this.notificationPage.providePublicEncryptionKey(this.extensionId)
+  }
+
+  async decrypt() {
+    if (!this.extensionId) {
+      throw NO_EXTENSION_ID_ERROR
+    }
+
+    await this.notificationPage.decryptMessage(this.extensionId)
+  }
+
   /// -------------------------------------------
   /// ---------- EXPERIMENTAL FEATURES ----------
   /// -------------------------------------------
@@ -429,12 +453,4 @@ export class MetaMask {
   async closeTransactionDetails() {
     await this.homePage.closeTransactionDetails()
   }
-
-  async addNewToken() {
-    if (!this.extensionId) {
-      throw NO_EXTENSION_ID_ERROR
-    }
-
-    await this.notificationPage.addNewToken(this.extensionId)
-  }
 }
diff --git a/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts b/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts
new file mode 100644
index 000000000..2edde6e5b
--- /dev/null
+++ b/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts
@@ -0,0 +1,10 @@
+import Selectors from '../selectors'
+import type { Page } from '@playwright/test'
+
+export async function providePublicEncryptionKey(notificationPage: Page) {
+  await notificationPage.locator(Selectors.ActionFooter.confirmActionButton).click()
+}
+
+export async function decryptMessage(notificationPage: Page) {
+  await notificationPage.locator(Selectors.ActionFooter.confirmActionButton).click()
+}
diff --git a/wallets/metamask/src/pages/NotificationPage/actions/index.ts b/wallets/metamask/src/pages/NotificationPage/actions/index.ts
index d54f0c031..3a2f10500 100644
--- a/wallets/metamask/src/pages/NotificationPage/actions/index.ts
+++ b/wallets/metamask/src/pages/NotificationPage/actions/index.ts
@@ -5,3 +5,4 @@ export * from './approvePermission'
 export * from './transaction'
 export * from './network'
 export * from './token'
+export * from './encryption'
diff --git a/wallets/metamask/src/pages/NotificationPage/page.ts b/wallets/metamask/src/pages/NotificationPage/page.ts
index 59e3bc13f..2a9e295b7 100644
--- a/wallets/metamask/src/pages/NotificationPage/page.ts
+++ b/wallets/metamask/src/pages/NotificationPage/page.ts
@@ -9,7 +9,9 @@ import {
   signSimpleMessage,
   signStructuredMessage,
   token,
-  transaction
+  transaction,
+  providePublicEncryptionKey,
+  decryptMessage
 } from './actions'
 import Selectors from './selectors'
 
@@ -139,4 +141,16 @@ export class NotificationPage {
 
     await token.addNew(notificationPage)
   }
+
+  async providePublicEncryptionKey(extensionId: string) {
+    const notificationPage = await getNotificationPageAndWaitForLoad(this.page.context(), extensionId)
+
+    await providePublicEncryptionKey(notificationPage)
+  }
+
+  async decryptMessage(extensionId: string) {
+    const notificationPage = await getNotificationPageAndWaitForLoad(this.page.context(), extensionId)
+
+    await decryptMessage(notificationPage)
+  }
 }
diff --git a/wallets/metamask/test/e2e/metamask/encrypt.spec.ts b/wallets/metamask/test/e2e/metamask/encrypt.spec.ts
new file mode 100644
index 000000000..78981df68
--- /dev/null
+++ b/wallets/metamask/test/e2e/metamask/encrypt.spec.ts
@@ -0,0 +1,34 @@
+import { testWithMetaMask } from '../testWithMetaMask'
+
+const test = testWithMetaMask
+
+const { expect } = test
+
+test('should provide public encryption key', async ({ page, metamask }) => {
+  await page.locator('#getEncryptionKeyButton').click()
+  await metamask.providePublicEncryptionKey()
+
+  await expect(page.locator('#encryptionKeyDisplay')).toHaveText('mtrHp1WHZM9rxF2Ilot9Hie5XmQcKCf7oDQ1DpGkTSI=')
+})
+
+test('should encrypt and decrypt a message', async ({ page, metamask }) => {
+  await page.locator('#getEncryptionKeyButton').click()
+  await metamask.providePublicEncryptionKey()
+  await expect(page.locator('#encryptionKeyDisplay')).toHaveText('mtrHp1WHZM9rxF2Ilot9Hie5XmQcKCf7oDQ1DpGkTSI=')
+
+  // `fill` does not trigger buttons validation, so we use `type` instead
+  await page.locator('#encryptMessageInput').type('Hello, world')
+
+  await page.locator('#encryptButton').click()
+
+  const encryptedMessage = await page.locator('#ciphertextDisplay').textContent()
+
+  expect(encryptedMessage).toContain('0x7b')
+
+  await page.locator('#decryptButton').click()
+  await metamask.decrypt()
+
+  const decryptedMessage = await page.locator('#cleartextDisplay').textContent()
+
+  expect(decryptedMessage).toEqual('Hello, world')
+})

From 6f5034be00385c5f67b8d8c3ab13c2ebd5aec103 Mon Sep 17 00:00:00 2001
From: matstyler <mat.franczuk@gmail.com>
Date: Wed, 28 Feb 2024 15:25:28 +0100
Subject: [PATCH 2/2] fix: cleanup

---
 .../src/pages/NotificationPage/actions/encryption.ts        | 2 +-
 wallets/metamask/src/pages/NotificationPage/page.ts         | 6 +++---
 .../metamask/src/pages/OnboardingPage/selectors/index.ts    | 4 ++--
 .../metamask/src/utils/getNotificationPageAndWaitForLoad.ts | 4 +++-
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts b/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts
index 2edde6e5b..365371806 100644
--- a/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts
+++ b/wallets/metamask/src/pages/NotificationPage/actions/encryption.ts
@@ -1,5 +1,5 @@
-import Selectors from '../selectors'
 import type { Page } from '@playwright/test'
+import Selectors from '../selectors'
 
 export async function providePublicEncryptionKey(notificationPage: Page) {
   await notificationPage.locator(Selectors.ActionFooter.confirmActionButton).click()
diff --git a/wallets/metamask/src/pages/NotificationPage/page.ts b/wallets/metamask/src/pages/NotificationPage/page.ts
index 2a9e295b7..ebed5cced 100644
--- a/wallets/metamask/src/pages/NotificationPage/page.ts
+++ b/wallets/metamask/src/pages/NotificationPage/page.ts
@@ -5,13 +5,13 @@ import {
   type GasSetting,
   approvePermission,
   connectToDapp,
+  decryptMessage,
   network,
+  providePublicEncryptionKey,
   signSimpleMessage,
   signStructuredMessage,
   token,
-  transaction,
-  providePublicEncryptionKey,
-  decryptMessage
+  transaction
 } from './actions'
 import Selectors from './selectors'
 
diff --git a/wallets/metamask/src/pages/OnboardingPage/selectors/index.ts b/wallets/metamask/src/pages/OnboardingPage/selectors/index.ts
index 3fcc0ceb4..c681569c8 100644
--- a/wallets/metamask/src/pages/OnboardingPage/selectors/index.ts
+++ b/wallets/metamask/src/pages/OnboardingPage/selectors/index.ts
@@ -21,5 +21,5 @@ export default {
   WalletCreationSuccessPageSelectors,
 
   // 5th Page
-  PinExtensionPageSelectors
-}
+  PinExtensionPageSelectors,
+};
diff --git a/wallets/metamask/src/utils/getNotificationPageAndWaitForLoad.ts b/wallets/metamask/src/utils/getNotificationPageAndWaitForLoad.ts
index 0f444598c..2a8951782 100644
--- a/wallets/metamask/src/utils/getNotificationPageAndWaitForLoad.ts
+++ b/wallets/metamask/src/utils/getNotificationPageAndWaitForLoad.ts
@@ -9,7 +9,9 @@ export async function getNotificationPageAndWaitForLoad(context: BrowserContext,
   let notificationPage = context.pages().find(isNotificationPage)
 
   if (!notificationPage) {
-    notificationPage = await context.waitForEvent('page', { predicate: isNotificationPage })
+    notificationPage = await context.waitForEvent('page', {
+      predicate: isNotificationPage
+    })
   }
 
   // Set pop-up window viewport size to resemble the actual MetaMask pop-up window.