feat(Browser extension): Firefox support (#74)

* chore(Browser extension): Improve Firefox support by moving fetch to background / service worker * chore(Browser extension): Improve Firefox support customizing manifest for MV2 * chore(Browser extension): Improve async fetch from background script * chore(Browser extension): Improve async fetch from background script (take 2) * chore(Browser extension): eslint fix * chore(Browser extension): Background script optimise for listener registration * chore(Browser extension): Dynamically pull config in Threat Composer inject script * chore(Browser extension): Added logo * chore(Browser extension): Revised logo * docs(Browser extension): README * build(Browser extension): Include Firefox build in build task * docs(Browser extension): Updated README for local install instructions --------- Co-authored-by: Darran Boyd <[email protected]>
awslabs · Jan 31, 2024 · 7ef06cd · 7ef06cd
1 parent 9c501cb
commit 7ef06cd
Show file tree

Hide file tree

Showing 10 changed files with 277 additions and 187 deletions.
diff --git a/.projenrc.ts b/.projenrc.ts
@@ -294,7 +294,7 @@ browserExtensionProject.addTask("dev:firefox", {
   exec: "wxt --browser firefox",
 });
 
-browserExtensionProject.compileTask.reset("wxt build");
+browserExtensionProject.compileTask.reset("wxt build -b chrome; wxt build -b firefox");
 
 browserExtensionProject.addTask("zip", {
   exec: "wxt zip",

diff --git a/README.md b/README.md
diff --git a/packages/threat-composer-app-browser-extension/.projen/tasks.json b/packages/threat-composer-app-browser-extension/.projen/tasks.json
diff --git a/packages/threat-composer-app-browser-extension/.wxt/types/paths.d.ts b/packages/threat-composer-app-browser-extension/.wxt/types/paths.d.ts
@@ -5,6 +5,7 @@ declare module "wxt/browser" {
   export type PublicPath =
     | "/background.js"
     | "/content-script.js"
+    | "/icon-128.png"
     | "/popup.html"
     | "/scriptInjectForCodeCatalyst.js"
     | "/scriptInjectForThreatComposer.js"

diff --git a/packages/threat-composer-app-browser-extension/README.md b/packages/threat-composer-app-browser-extension/README.md
@@ -1 +1,60 @@
-# replace this
+# Threat Composer Browser Extension
+
+## Introduction
+
+This extension/add-on allows you to view web accessible Threat Composer exports (.tc.json) with one click directly in your browser. The Threat Composer web-app is baked into the extension itself.
+
+The extension supports the following integrations where a ‘View in Threat Composer’ button is added to the page:
+
+- GitHub Code Browser
+- Amazon Code Catalyst
+- ‘View Raw’ anywhere online so long as the URL includes reference to `.tc.json` (Note: on Firefox it does not work on `githubusercontent.com` due to the `sandbox` CSP directive)
+
+## Build
+
+1. Clone this repo
+1. Run the main build script (`./scripts/build.sh`) to build everything (including the browser extension)
+
+For any other guidance see the [Development](../../README.md#development) section of the main [README](../../README.md)
+
+## Load locally
+
+### Google Chrome
+
+1. Open Chrome, then goto `chrome://extensions`
+1. Enable 'Developer mode'
+1. Click on 'Load unpacked'
+1. Point it at the `./.output/chrome-mv3` directory and click Open
+
+### Mozilla Firefox
+
+1. Open Firefox, then goto `about:debugging`
+2. Click on 'This Firefox' then 'Load Temporary Add-on...'
+3. Point it at any file with the `./.output/firefox-mv2` and click Open
+
+## Create ZIP file
+
+### Google Chrome
+
+1. Go to the root of the extension package - `cd ./packages/threat-composer-app-browser-extension`
+1. Run `yarn run zip` - look in `./.output/` for ZIP file
+1. To load locally open Chrome, then goto `chrome://extensions`, enable developer mode, then drag-and-drop the ZIP file onto the page to load.
+
+### Mozilla Firefox
+
+1. Go to the root of the extension package - `cd ./packages/threat-composer-app-browser-extension`
+1. Run `yarn run zip:firefox` - look in `./.output/` for ZIP file
+
+## Development
+
+### Google Chrome
+
+1. Go to the root of the extension package - `cd ./packages/threat-composer-app-browser-extension`
+1. Run `yarn run dev`
+1. In your browser navigate to a hosted Threat Composer file on a supported integration - [example1](https://github.com/awslabs/threat-composer/blob/main/packages/threat-composer/src/data/workspaceExamples/ThreatComposer.tc.json), [example2](https://github.com/awslabs/threat-composer/blob/main/packages/threat-composer/src/data/workspaceExamples/GenAIChatbot.tc.json) and [example3](https://raw.githubusercontent.com/awslabs/threat-composer/main/packages/threat-composer/src/data/workspaceExamples/GenAIChatbot.tc.json)
+
+### Mozilla Firefox
+
+1. Go to the root of the extension package - `cd ./packages/threat-composer-app-browser-extension`
+1. Run `yarn run dev:firefox`
+1. In your browser navigate to a hosted Threat Composer file on a supported integration - [example1](https://github.com/awslabs/threat-composer/blob/main/packages/threat-composer/src/data/workspaceExamples/ThreatComposer.tc.json) and [example2](https://github.com/awslabs/threat-composer/blob/main/packages/threat-composer/src/data/workspaceExamples/GenAIChatbot.tc.json). Note: on Firefox it does not work on `githubusercontent.com` due to the `sandbox` CSP directive.
diff --git a/packages/threat-composer-app-browser-extension/src/entrypoints/background.ts b/packages/threat-composer-app-browser-extension/src/entrypoints/background.ts
@@ -14,14 +14,15 @@
   limitations under the License.
  ******************************************************************************************************************** */
 
-import { getExtensionConfig, ThreatComposerTarget } from './popup/config';
+import { ThreatComposerTarget, getExtensionConfig } from './popup/config';
 import { logDebugMessage } from '../debugLogger';
 
+
 export default defineBackground(() => {
-  getExtensionConfig().then(config => {
-    logDebugMessage(config, 'index.hml is here:' + browser.runtime.getURL('index.html'));
 
-    browser.runtime.onMessage.addListener(function (request: any) {
+  browser.runtime.onMessage.addListener(function (request: any, sender: any, sendResponse: any) {
+
+    getExtensionConfig().then(config => {
       const tcViewer = config.target;
       let tcUrlCreate = '';
       let tcUrlUpdate = '';
@@ -37,8 +38,7 @@ export default defineBackground(() => {
         tcUrlUpdate = tcUrlCreate;
       }
 
-      if (request.schema) {
-        //This is likely the JSON from a threat model
+      if (request.schema) { //This is likely the JSON from a threat model
         logDebugMessage(config, 'Message recieved - Threat Model JSON');
 
         browser.storage.local.set({ threatModel: request }).then(() => {
@@ -52,9 +52,19 @@ export default defineBackground(() => {
             browser.tabs.create({ url: tcUrlCreate });
           }
         });
+        sendResponse({});
+      } else if (request.url) { //This is likely the a request to proxy a Fetch
+        sendResponse(fetch(request.url)
+          .then((response) => response.json())
+          .catch((error) => {
+            logDebugMessage({ debug: true } as any, error);
+          }));
       }
+    }).catch((error) => {
+      logDebugMessage({ debug: true } as any, error);
     });
-  }).catch((error) => {
-    logDebugMessage({ debug: true } as any, error);
+
+    // As we will reply asynchronously to the request, we need to tell chrome to wait for our response
+    return true;
   });
 });
diff --git a/packages/threat-composer-app-browser-extension/src/entrypoints/content-script.ts b/packages/threat-composer-app-browser-extension/src/entrypoints/content-script.ts
@@ -33,50 +33,54 @@ interface TCCodeBrowserState {
   stopProcessing: boolean;
 }
 
+function forwardFetchToBackground(message: any): Promise<TCJSONSimplifiedSchema> {
+  return browser.runtime.sendMessage(message).then((response: any) => {
+    if (!response) throw browser.runtime.lastError;
+    return response;
+  });
+}
 
 function isLikelyThreatComposerSchema(JSONobj: TCJSONSimplifiedSchema) {
   return JSONobj.schema ? true : false;
 };
 
 async function getTCJSONCandidate(url: string, element: HTMLElement, config: TCConfig) {
-  const tcJSONCandidate: TCJSONSimplifiedSchema = await fetch(url)
-    .then(function (response) {
+  return forwardFetchToBackground({ url: url })
+    .then(function (tcJSONCandidate) {
       logDebugMessage(config, 'Able to get a JSON candidate');
-      return response.json();
+      if (tcJSONCandidate && isLikelyThreatComposerSchema(tcJSONCandidate)) {
+        logDebugMessage(config,
+          'Looks like it could be a Threat Composer file, enabling ' +
+          element.textContent +
+          ' button',
+        );
+        element.onclick = function () {
+          logDebugMessage(config,
+            'Sending message with candicate JSON object back service worker / background script',
+          );
+          browser.runtime.sendMessage(tcJSONCandidate);
+        };
+
+        switch (element.tagName) { //"Enable" button / anchor
+          case 'BUTTON':
+            (element as HTMLInputElement).disabled = false;
+            break;
+          case 'A': //Anchor
+            element.style.pointerEvents = 'auto';
+            break;
+        }
+
+      } else {
+        logDebugMessage(config,
+          "Does NOT look like it's a Threat Composer file, NOT enabling " +
+          element.textContent +
+          ' button',
+        );
+      }
     })
     .catch(function (error) {
       logDebugMessage(config, 'Error during fetch: ' + error.message);
     });
-
-  if (tcJSONCandidate && isLikelyThreatComposerSchema(tcJSONCandidate)) {
-    logDebugMessage(config,
-      'Looks like it could be a Threat Composer file, enabling ' +
-      element.textContent +
-      ' button',
-    );
-    element.onclick = function () {
-      logDebugMessage(config,
-        'Sending message with candicate JSON object back service worker / background script',
-      );
-      browser.runtime.sendMessage(tcJSONCandidate);
-    };
-
-    switch (element.tagName) { //"Enable" button / anchor
-      case 'BUTTON':
-        (element as HTMLInputElement).disabled = false;
-        break;
-      case 'A': //Anchor
-        element.style.pointerEvents = 'auto';
-        break;
-    }
-
-  } else {
-    logDebugMessage(config,
-      "Does NOT look like it's a Threat Composer file, NOT enabling " +
-      element.textContent +
-      ' button',
-    );
-  }
 };
 
 async function handleRawFile(config: TCConfig) {

diff --git a/packages/threat-composer-app-browser-extension/src/public/icon-128.png b/packages/threat-composer-app-browser-extension/src/public/icon-128.png
diff --git a/packages/threat-composer-app-browser-extension/src/public/scriptInjectForThreatComposer.js b/packages/threat-composer-app-browser-extension/src/public/scriptInjectForThreatComposer.js
@@ -1,64 +1,76 @@
 const browser = window.browser ? window.browser : window.chrome;
 
-const debug = true; //TODO: Get this value from config.
+async function getConfig() {
+  return browser.storage.local.get(["tcConfig"]).then((config) => {
+    if (config.tcConfig) {
+      return config.tcConfig;
+    } else {
+      return { debug: true };
+    }
+  });
+}
 
-function logDebugMessage(msg) {
+async function logDebugMessage(msg) {
+  config = await getConfig();
   const debugPrefix = "ThreatComposerExtension: ";
-  if (debug) console.log(debugPrefix + msg);
+  if (config.debug) console.log(debugPrefix + msg);
 }
 
 function loadThreatModel(source) {
-  logDebugMessage("loadThreatModel: triggered" + " source:" + source);
   if (window.threatcomposer.setCurrentWorkspaceData) {
     logDebugMessage(
-      "loadThreatModel: window.threatcomposer.setCurrentWorkspaceData is NOT undefined"
+      "Attempting to retrieve threat model from local storage" +
+        " source: " +
+        source
     );
     browser.storage.local
       .get("threatModel")
       .then((result) => {
         const answer = result.threatModel;
         if (answer != undefined) {
           if (answer.schema) {
-            logDebugMessage(answer);
             logDebugMessage(
-              "loadThreatModel: got valid JSON loading threat model, and disconnecting observer" +
-                " source:" +
+              "Got valid JSON loading threat model, triggering to load into current workspace." +
+                " source: " +
                 source
             );
             window.threatcomposer.setCurrentWorkspaceData(answer);
           } else {
             logDebugMessage(
-              "loadThreatModel: Did NOT get valid JSON loading threat model, NOT disconnecting observer" +
-                " source:" +
+              "Did NOT get valid JSON. Will NOT trigger loading of threat model." +
+                " source: " +
                 source
             );
           }
         } else {
           logDebugMessage(
-            "loadThreatModel: Did NOT get JSON object, NOT disconnecting observer" +
-              " source:" +
+            "Did NOT get JSON object. Will NOT trigger loading of threat model." +
+              " source: " +
               source
           );
         }
       })
       .catch((error) => {
-        console.log(debugPrefix + error);
+        logDebugMessage(
+          "Error during when attempting to retrieve threat model from local storage: " +
+            error.message
+        );
       });
   } else {
     logDebugMessage(
-      "loadThreatModel: window.threatcomposer.setCurrentWorkspaceData is undefined" +
-        " source:" +
+      "window.threatcomposer.setCurrentWorkspaceData is undefined. Cannot proceed until it is" +
+        " source: " +
         source
     );
   }
 }
 
 logDebugMessage(
-  "Adding visibility event listener and mutation observer to trigger load of threat model"
+  "Adding DOMContentLoaded and visibility event listener to trigger load of threat model"
 );
 
 document.addEventListener("visibilitychange", (event) => {
-  loadThreatModel("visibility change");
+  loadThreatModel("visibilitychange");
 });
 
 async function sleep(ms) {
@@ -70,9 +82,9 @@ async function sleep(ms) {
 document.addEventListener("DOMContentLoaded", async (event) => {
   while (!window.threatcomposer.setCurrentWorkspaceData) {
     logDebugMessage(
-      "Waiting for window.threatcomposer.setCurrentWorkspaceData"
+      "Waiting for window.threatcomposer.setCurrentWorkspaceData to be ready."
     );
     await sleep(50);
   }
-  loadThreatModel("onload");
+  loadThreatModel("DOMContentLoaded");
 });