From 6b87f7add3e7590da883c4a2631c844d6899e6e9 Mon Sep 17 00:00:00 2001
From: Kenneth Hendricks <50819541+kenhendricks00@users.noreply.github.com>
Date: Sun, 20 Oct 2024 23:54:16 -0400
Subject: [PATCH] Update index.js

Updated the proper URL, if the whole site is labeled or if a specific page
---
 pub/index.js | 75 +++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 15 deletions(-)

diff --git a/pub/index.js b/pub/index.js
index 75e9bcc..42fd595 100644
--- a/pub/index.js
+++ b/pub/index.js
@@ -8,26 +8,56 @@ document.addEventListener("DOMContentLoaded", async () => {
   // Helper function to normalize URLs (removes trailing slashes)
   const normalizeUrl = (url) => url.replace(/\/+$/, "").trim();
 
+  // Helper function to extract the root domain from a URL
+  function extractRootDomain(url) {
+    let urlObj = new URL(url);
+    return `${urlObj.protocol}//${urlObj.hostname}`;
+  }
+
   try {
     // Get the active tab's URL
-    const [activeTab] = await browser.tabs.query({ active: true, currentWindow: true });
+    const [activeTab] = await browser.tabs.query({
+      active: true,
+      currentWindow: true,
+    });
 
     if (!activeTab || !activeTab.url) {
       throw new Error("No active tab found or URL is unavailable.");
     }
 
     const currentUrl = normalizeUrl(activeTab.url);
-    console.log(`Active tab URL: ${currentUrl}`);
+    const rootDomain = extractRootDomain(currentUrl);
+    console.log(`Active tab URL: ${currentUrl}, Root domain: ${rootDomain}`);
 
     // Send a message to the background script to check the site's status
-    const response = await browser.runtime.sendMessage({ action: "checkSiteStatus", url: currentUrl });
+    const response = await browser.runtime.sendMessage({
+      action: "checkSiteStatus",
+      url: currentUrl,
+    });
 
     if (!response || !response.status) {
-      throw new Error("Failed to retrieve site status from the background script.");
+      throw new Error(
+        "Failed to retrieve site status from the background script."
+      );
     }
 
+    // Determine if the root domain or the full URL is marked as safe/starred
+    const isRootDomainMarked = await browser.runtime.sendMessage({
+      action: "checkSiteStatus",
+      url: rootDomain,
+    });
+
     // Handle different site statuses and update the UI accordingly
-    handleStatusUpdate(response.status, currentUrl);
+    if (
+      isRootDomainMarked.status === response.status &&
+      response.status !== "no_data"
+    ) {
+      // If both the root domain and the current URL share the same status, show the root domain in the message
+      handleStatusUpdate(response.status, rootDomain);
+    } else {
+      // Otherwise, show the current URL in the message
+      handleStatusUpdate(response.status, currentUrl);
+    }
   } catch (error) {
     console.error("Error while checking site status:", error);
     errorMessage.textContent = `Error: ${error.message}`;
@@ -37,27 +67,42 @@ document.addEventListener("DOMContentLoaded", async () => {
   /**
    * Updates the UI based on the site status
    * @param {string} status - The status of the site (e.g., "safe", "unsafe")
-   * @param {string} url - The URL of the current tab
+   * @param {string} displayUrl - The URL or root domain to display in the message
    */
-  function handleStatusUpdate(status, url) {
+  function handleStatusUpdate(status, displayUrl) {
     switch (status) {
       case "unsafe":
-        updateUI("unsafe", `${url} is flagged as <strong>unsafe</strong>. Be cautious when interacting with this site.`);
+        updateUI(
+          "unsafe",
+          `${displayUrl} is flagged as <strong>unsafe</strong>. Be cautious when interacting with this site.`
+        );
         break;
       case "potentially_unsafe":
-        updateUI("potentially_unsafe", `${url} is <strong>potentially unsafe</strong>. Proceed with caution.`);
+        updateUI(
+          "potentially_unsafe",
+          `${displayUrl} is <strong>potentially unsafe</strong>. Proceed with caution.`
+        );
         break;
       case "safe":
-        updateUI("safe", `${url} is <strong>safe</strong> to browse.`);
+        updateUI("safe", `${displayUrl} is <strong>safe</strong> to browse.`);
         break;
       case "starred":
-        updateUI("starred", `${url} is a <strong>starred</strong> site.`);
+        updateUI(
+          "starred",
+          `${displayUrl} is a <strong>starred</strong> site.`
+        );
         break;
       case "no_data":
-        updateUI("no_data", `No data available for <strong>${url}</strong>.`);
+        updateUI(
+          "no_data",
+          `No data available for <strong>${displayUrl}</strong>.`
+        );
         break;
       default:
-        updateUI("unknown", `An unknown status was received for <strong>${url}</strong>.`);
+        updateUI(
+          "unknown",
+          `An unknown status was received for <strong>${displayUrl}</strong>.`
+        );
     }
   }
 
@@ -74,13 +119,13 @@ document.addEventListener("DOMContentLoaded", async () => {
       starred: "../res/icons/starred.png",
       no_data: "../res/ext_icon_144.png",
       error: "../res/icons/error.png",
-      unknown: "../res/ext_icon_144.png"
+      unknown: "../res/ext_icon_144.png",
     };
 
     // Update the icon and message
     statusIcon.src = icons[status] || icons["unknown"];
     statusMessage.innerHTML = message || "An unknown error occurred.";
-    
+
     // Add a small animation when the status changes
     statusIcon.classList.add("active");
     setTimeout(() => statusIcon.classList.remove("active"), 300);