Implement one-time code auto-fill and copy

src/background/content-script.ts

@@ -92,6 +92,17 @@ const observe = (input: HTMLInputElement, form: LoginForm) => {
           warnings,
         };
       }
+
+      case "FILL_ONE_TIME_CODE": {
+        const { username, code } = message;
+        const warnings = fillLoginForm(form, username, code);
+        destroy();
+
+        return {
+          success: true,
+          warnings,
+        };
+      }
     }
   };
 

src/background/fill-one-time-code.ts

@@ -0,0 +1,31 @@
+import { fillLoginForm, getLoginForms } from "../utils/dom";
+
+declare global {
+  interface Window {
+    iCPFillOneTimeCode: (
+      username: string,
+      code: string,
+    ) => Promise<{ success: true; warnings: string[] }>;
+  }
+}
+
+window.iCPFillOneTimeCode = async (username, code) => {
+  const warnings: string[] = [];
+
+  const forms = getLoginForms();
+
+  if (forms.length === 0) {
+    throw new Error("AutoFill failed: no one-time code field on page");
+  } else if (forms.length > 1) {
+    warnings.push(
+      "Multiple one-time codes detected on page, only filling the last",
+    );
+  }
+
+  warnings.push(...fillLoginForm(forms[forms.length - 1], username, code));
+
+  return {
+    success: true,
+    warnings,
+  };
+};

src/background/index.ts

@@ -145,6 +145,86 @@ browser.runtime.onMessage.addListener(async (message, sender) => {
             message.url,
           ),
         };
+
+      case "FETCH_ONE_TIME_CODE":
+        return {
+          success: true,
+          oneTimeCode: await getAPI().fetchOneTimeCode(
+            message.tabId,
+            message.url,
+            message.oneTimeCode.username,
+          ),
+        };
+
+      case "FILL_ONE_TIME_CODE": {
+        const { username, code } = await getAPI().fetchOneTimeCode(
+          message.tabId ?? sender.tab?.id ?? -1,
+          message.url,
+          message.loginName,
+        );
+
+        const tab = sender.tab ?? (await browser.tabs.get(message.tabId));
+        if (tab.id === undefined)
+          throw new Error("AutoFill failed: tab no longer exists");
+        if (!tab.active)
+          throw new Error("AutoFill failed: tab is no longer active");
+        if (tab.url !== message.url)
+          throw new Error("AutoFill failed: tab has changed URL");
+
+        if (message.forwardToContentScript) {
+          const { success, warnings } = await browser.tabs.sendMessage(tab.id, {
+            cmd: "FILL_ONE_TIME_CODE",
+            username,
+            code,
+          });
+
+          (warnings as string[]).forEach((warning) => console.warn(warning));
+
+          return {
+            success,
+            warnings,
+          };
+        } else {
+          await browser.scripting.executeScript({
+            target: {
+              tabId: tab.id,
+            },
+            files: ["./fill_one_time_code.js"],
+          });
+
+          const [{ result, error }] = await browser.scripting.executeScript({
+            target: {
+              tabId: tab.id,
+            },
+            func: (username, code) => window.iCPFillOneTimeCode(username, code),
+            args: [username, code],
+          });
+
+          if (error !== undefined) throw error;
+
+          const { success, warnings } = result;
+          (warnings as string[]).forEach((warning) => console.warn(warning));
+
+          return {
+            success,
+            warnings,
+          };
+        }
+      }
+
+      case "COPY_ONE_TIME_CODE": {
+        const { code } = await getAPI().fetchOneTimeCode(
+          message.tabId,
+          message.url,
+          message.loginName,
+        );
+
+        await navigator.clipboard.writeText(code);
+
+        return {
+          success: true,
+        };
+      }
     }
   } catch (e) {
     console.error(e);

src/ui/popup/one-time-codes.tsx

@@ -1,11 +1,16 @@
-import styles from "./one-time-codes.module.scss";
+import { useEffect, useState } from "react";
 import { useNavigate, useOutletContext } from "react-router-dom";
-import { useEffect } from "react";
+import browser from "webextension-polyfill";
 import { HomeContext } from "./home";
+import { OneTimeCode } from "../shared/hooks/use-one-time-codes";
+import { ErrorView } from "../shared/error";
+import { CopyIcon } from "../shared/icons/copy";
+import styles from "./one-time-codes.module.scss";
 
 export function OneTimeCodesView() {
   const navigate = useNavigate();
-  const { oneTimeCodes } = useOutletContext<HomeContext>();
+  const { tab, oneTimeCodes } = useOutletContext<HomeContext>();
+  const [fillError, setFillError] = useState<string>();
 
   const hasOneTimeCode = oneTimeCodes.length !== 0;
 
@@ -14,37 +19,65 @@ export function OneTimeCodesView() {
     navigate("/");
   }, [hasOneTimeCode]);
 
+  const handleFillOneTimeCode = async (
+    oneTimeCode: OneTimeCode,
+    action: "FILL" | "COPY" = "FILL",
+  ) => {
+    if (tab?.id === undefined || tab?.url === undefined) return;
+
+    setFillError(undefined);
+
+    try {
+      // Can't use FETCH_ONE_TIME_CODE here
+      // See https://bugzilla.mozilla.org/show_bug.cgi?id=1292701
+      const { success, error } = await browser.runtime.sendMessage({
+        cmd: `${action}_ONE_TIME_CODE`,
+        tabId: tab.id,
+        url: tab.url,
+        oneTimeCode,
+      });
+
+      if (error !== undefined || !success) throw error;
+
+      window.close();
+    } catch (e: any) {
+      setFillError(e.message ?? e.toString());
+    }
+  };
+
+  if (fillError !== undefined) return <ErrorView error={fillError} />;
+
+  if (!hasOneTimeCode) return null;
+
   return (
     <div className={styles.oneTimeCodes}>
-      {oneTimeCodes.length > 0 ? (
-        <>
-          <h2>Select a verification code to use:</h2>
-
-          <ul>
-            {oneTimeCodes.map((oneTimeCode, i) => (
-              <li
-                key={i}
-                onClick={(e) => {
-                  e.preventDefault();
-                  // TODO
-                }}
-              >
-                <div>
-                  <span>{oneTimeCode.username}</span>
-                  <span>{oneTimeCode.domain}</span>
-                </div>
-              </li>
-            ))}
-          </ul>
-        </>
-      ) : (
-        <p>
-          <br />
-          No one time codes saved on this website.
-          <br />
-          <br />
-        </p>
-      )}
+      <h2>Select a verification code to use:</h2>
+
+      <ul>
+        {oneTimeCodes.map((oneTimeCode, i) => (
+          <li
+            key={i}
+            onClick={(e) => {
+              e.preventDefault();
+              handleFillOneTimeCode(oneTimeCode);
+            }}
+          >
+            <div>
+              <span>{oneTimeCode.username}</span>
+              <span>{oneTimeCode.domain}</span>
+            </div>
+            <a
+              onClick={(e) => {
+                e.preventDefault();
+                e.stopPropagation();
+                handleFillOneTimeCode(oneTimeCode, "COPY");
+              }}
+            >
+              <CopyIcon title="Copy one-time code to clipboard" />
+            </a>
+          </li>
+        ))}
+      </ul>
     </div>
   );
 }

src/ui/shared/error.tsx

@@ -15,6 +15,13 @@ export function ErrorView({ error }: Props) {
           the login form open.
         </>
       );
+    } else if (error === "AutoFill failed: no one-time code field on page") {
+      return (
+        <>
+          No one-time code field was found on the current page. Make sure you
+          have the login form open.
+        </>
+      );
     } else if (error === "AutoFill failed: tab is no longer active") {
       return (
         <>

src/utils/api.ts

@@ -443,6 +443,68 @@ export class ApplePasswordManager extends EventEmitter {
     }
   }
 
+  async fetchOneTimeCode(tabId: number, url: string, username: string) {
+    if (this.session === undefined)
+      throw new Error("Invalid session state: not initialized");
+
+    const sdata = this.session.serialize(
+      await this.session.encrypt({
+        ACT: Action.SEARCH,
+        TYPE: "oneTimeCodes",
+        frameURLs: [url],
+        username,
+      }),
+    );
+
+    const { payload } = await this._postMessage(
+      Command.DID_FILL_ONE_TIME_CODE,
+      {
+        tabId,
+        frameId: 0,
+        payload: {
+          QID: "CmdDidFillOneTimeCode",
+          SMSG: JSON.stringify({
+            TID: this.session.username,
+            SDATA: sdata,
+          }),
+        },
+      },
+      null,
+    );
+
+    // macOS sends this as an object, Windows as a string
+    if (typeof payload.SMSG === "string")
+      payload.SMSG = JSON.parse(payload.SMSG);
+
+    if (payload.SMSG.TID !== this.session.username)
+      throw new Error("Invalid server response: destined to another session");
+
+    let response;
+    try {
+      const data = await this.session.decrypt(
+        this.session.deserialize(payload.SMSG.SDATA),
+      );
+      response = JSON.parse(data.toString("utf8"));
+    } catch (e) {
+      throw new Error("Invalid server response: missing payload");
+    }
+
+    switch (response.STATUS) {
+      case QueryStatus.SUCCESS: {
+        const { username, domain, source, code } = response.Entries[0];
+        return {
+          username,
+          domain,
+          source,
+          code,
+        };
+      }
+
+      default:
+        throwQueryStatusError(response.STATUS);
+    }
+  }
+
   async close() {
     const port = this.port;
     if (port === undefined) return;