From 9895845cb53e19e872449cff036f2f2c4cf42323 Mon Sep 17 00:00:00 2001
From: Luca Scalzotto <luca@scalzotto.nl>
Date: Fri, 28 Jun 2024 11:35:03 +0200
Subject: [PATCH] Load content script in iframe, working again

---
 manifest.json               |  5 +++--
 src/TimeChimpApi.ts         | 26 +++++++++++++++++++-------
 src/background/timechimp.ts | 17 ++++++++++++++++-
 src/content/storage.ts      | 13 -------------
 4 files changed, 38 insertions(+), 23 deletions(-)
 delete mode 100644 src/content/storage.ts

diff --git a/manifest.json b/manifest.json
index 9baaa70..9547ee8 100644
--- a/manifest.json
+++ b/manifest.json
@@ -13,14 +13,15 @@
   "content_scripts": [
     {
       "matches": [
-        "https://app.timechimp.com/*"
+        "https://angular.timechimp.com/*"
       ],
       "js": [
         "build/content/index.js"
       ],
       "css": [
         "build/content/index.css"
-      ]
+      ],
+      "all_frames": true
     }
   ],
   "background": {
diff --git a/src/TimeChimpApi.ts b/src/TimeChimpApi.ts
index 78cdebc..3218d5c 100644
--- a/src/TimeChimpApi.ts
+++ b/src/TimeChimpApi.ts
@@ -1,12 +1,6 @@
-import { getTokenFromStorage } from './content/storage';
-
 export class TimeChimpApi {
     private async doFetch<T>(path: string): Promise<T> {
-        const token = getTokenFromStorage();
-        if (!token) {
-            throw new Error('No token found in local storage');
-        }
-
+        const token = this.getToken();
         const url = `https://web.timechimp.com${path}`;
         const response = await fetch(url, {
             headers: {
@@ -22,6 +16,24 @@ export class TimeChimpApi {
         return JSON.parse(body);
     }
 
+    private getToken(): string {
+        let token = window.localStorage.getItem('tc_auth_token');
+        if (!token) {
+            throw new Error('No token found in local storage');
+        }
+
+        // The token starts and ends with a double quote, remove that.
+        // Though also account for the fact that this could change.
+        if (token.startsWith('"')) {
+            token = token.substring(1);
+        }
+        if (token.endsWith('"')) {
+            token = token.substring(0, token.length - 1);
+        }
+
+        return token;
+    }
+
     public getCurrentUser(): Promise<User> {
         return this.doFetch('/api/user/current');
     }
diff --git a/src/background/timechimp.ts b/src/background/timechimp.ts
index c278745..7dda4d3 100644
--- a/src/background/timechimp.ts
+++ b/src/background/timechimp.ts
@@ -1,5 +1,20 @@
 import { Message } from '../message';
 
+/**
+ * The time registration form is an iframe in the webpage,
+ * so we need to wait for that to load, and store the frame id per tab.
+ * While unlikely, it is possible that someone has multiple TimeChimp tabs open.
+ */
+const frameByTab: Record<number, number> = {};
+chrome.webRequest.onCompleted.addListener(
+    (details) => {
+        if (details.type === 'sub_frame') {
+            frameByTab[details.tabId] = details.frameId;
+        }
+    },
+    { urls: ['https://angular.timechimp.com/*'] },
+);
+
 /**
  * Detects that the week has changed and the billability should be recalculated, on basis of requests to TimeChimp.
  */
@@ -40,7 +55,7 @@ chrome.webRequest.onCompleted.addListener(
 
 async function sendMessage(tabId: number, msg: Message) {
     await chrome.tabs
-        .sendMessage(tabId, msg)
+        .sendMessage(tabId, msg, { frameId: frameByTab[tabId] })
         .catch(() =>
             console.debug(
                 'Failed to send message to tab, content script is likely not loaded yet.',
diff --git a/src/content/storage.ts b/src/content/storage.ts
deleted file mode 100644
index 0c75390..0000000
--- a/src/content/storage.ts
+++ /dev/null
@@ -1,13 +0,0 @@
-export function getTokenFromStorage(): string | null {
-    for (let i = 0; i < window.localStorage.length; i++) {
-        const key = window.localStorage.key(i);
-        if (
-            key &&
-            key.startsWith('CognitoIdentityServiceProvider.') &&
-            key.endsWith('.accessToken')
-        ) {
-            return window.localStorage.getItem(key);
-        }
-    }
-    return null;
-}