From 12cdec7283efd08a46b30860f86a196c5d05c5ef Mon Sep 17 00:00:00 2001
From: Jose Luis <gskitcat@gmail.com>
Date: Thu, 28 Mar 2024 02:04:00 -0600
Subject: [PATCH] fix(validate url): make sure the extension only works on the
 Telegram K version

---
 public/manifest.json                  |  2 -
 src/App.tsx                           | 25 +++++++++++
 src/background.ts                     | 23 ++++++++--
 src/content_script.ts                 | 60 +++++++++++++++------------
 src/customHooks/useAddOn/index.ts     | 33 ++++++++++++++-
 src/helpers/content_script/types.ts   |  7 +++-
 src/structure/Gallery/index.tsx       |  5 ++-
 src/structure/configuration/index.tsx |  5 ++-
 8 files changed, 121 insertions(+), 39 deletions(-)

diff --git a/public/manifest.json b/public/manifest.json
index 046441c..f32ccde 100644
--- a/public/manifest.json
+++ b/public/manifest.json
@@ -11,8 +11,6 @@
     "16": "16x16.png",
     "256": "256x256.png"
   },
-  "background_color": "#78c2ad",
-
   "content_scripts": [
     {
       "matches": ["https://*.web.telegram.org/*"],
diff --git a/src/App.tsx b/src/App.tsx
index 08489d9..7874c78 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -6,6 +6,31 @@ import Gallery from "./structure/Gallery";
 function App() {
   const hook = useAddOn();
 
+  if (hook.isTelegramK === false)
+    return (
+      <main className={ui.addOn}>
+        <Header />
+        <p className="p-5">
+          In order to use this extension, you must use web telegram with the
+          following adress:
+          <a className="mx-2" target="_blank" href="https://web.telegram.org/k">
+            https://web.telegram.org/k
+          </a>
+        </p>
+        <p className="mx-5">
+          Once you enter the link, close and open this popup
+        </p>
+      </main>
+    );
+
+  if (hook.isTelegramK === null)
+    return (
+      <main className={ui.addOn}>
+        <Header />
+        <p className="p-5">Hold up a second, the extension is loading...</p>
+      </main>
+    );
+
   return (
     <main className={ui.addOn}>
       <Header />
diff --git a/src/background.ts b/src/background.ts
index 7a0308a..1ea6904 100644
--- a/src/background.ts
+++ b/src/background.ts
@@ -5,8 +5,8 @@ import { MessageBrowserActions } from "./helpers/content_script/types";
 // Listen for messages from the content script
 browser.runtime.onMessage.addListener(async (message: string) => {
   const parsedMessage: AddOnMessage = JSON.parse(message);
-  
-  console.log('background',parsedMessage);
+
+  console.log("background", parsedMessage);
 
   switch (parsedMessage.action) {
     case "crawledContent": {
@@ -43,5 +43,22 @@ async function requestFetchedImages() {
 async function receiveCrawledData(
   information: MessageBrowserActions<"crawledContent">
 ) {
-  browser.runtime.sendMessage(JSON.stringify(information));
+
+  const tab = await getTab();
+
+  browser.runtime.sendMessage(
+    JSON.stringify({
+      images: information,
+      urlTab: tab.url,
+    })
+  );
 }
+
+  async function getTab() {
+    const [tab] = await browser.tabs.query({
+      active: true,
+      currentWindow: true,
+    });
+
+    return tab;
+  }
diff --git a/src/content_script.ts b/src/content_script.ts
index b8c96ef..8385dcc 100644
--- a/src/content_script.ts
+++ b/src/content_script.ts
@@ -3,6 +3,7 @@ import { showHiddenMediaBtns } from "./helpers/content_script";
 import { MessageBrowserActions } from "./helpers/content_script/types";
 import { CrawledImageDom } from "./typesContentScript";
 
+// TODO: Instead of sending the content, render a button to download from the DOM
 // // Create a new instance of MutationObserver
 // const observer = new MutationObserver(async (mutationsList) => {
 //   for (const mutation of mutationsList) {
@@ -21,7 +22,7 @@ import { CrawledImageDom } from "./typesContentScript";
 //             blob: base64Image,
 //             height: targetElement.naturalWidth,
 //             width: targetElement.naturalHeight,
-            
+
 //           }],
 //         };
 
@@ -64,9 +65,8 @@ function hasSiblingButtons(element: Element) {
   return false;
 }
 
-function sendImages(information:MessageBrowserActions<"crawledContent">){
-
-  console.log(information)
+function sendImages(information: MessageBrowserActions<"crawledContent">) {
+  console.log(information);
 
   browser.runtime.sendMessage(JSON.stringify(information));
 }
@@ -99,10 +99,10 @@ async function downloadImageAsBase64(url: string | null): Promise<string> {
  * Get the images of the telegram chat
  * @returns Crawled images
  */
-async function getAllImages():Promise<CrawledImageDom[]> {
+async function getAllImages(): Promise<CrawledImageDom[]> {
   const container = document.getElementById("column-center");
 
-  if (container === null) return []
+  if (container === null) return [];
 
   // Use querySelectorAll to select all elements with class 'media-photo' inside the container
   const mediaPhotos = container.querySelectorAll(".media-photo");
@@ -110,39 +110,45 @@ async function getAllImages():Promise<CrawledImageDom[]> {
   // Convert the NodeList to an array (optional)
   const mediaPhotosArray = Array.from(mediaPhotos);
 
-  const photos = [...mediaPhotosArray].map(item=>(async function(){
-    if (hasSiblingButtons(item)) return null;
+  const photos = [...mediaPhotosArray].map((item) =>
+    (async function () {
+      if (hasSiblingButtons(item)) return null;
 
-    const imageElement = item as HTMLImageElement;
+      const imageElement = item as HTMLImageElement;
 
-    const base64Image = await downloadImageAsBase64(item.getAttribute("src"));
+      const base64Image = await downloadImageAsBase64(item.getAttribute("src"));
 
-    const crawledContent: CrawledImageDom = {
-      blob: base64Image,
-      height: imageElement.naturalWidth,
-      width: imageElement.naturalHeight,
-    };  
-  
-    return crawledContent
-  
-  })());
+      const crawledContent: CrawledImageDom = {
+        blob: base64Image,
+        height: imageElement.naturalWidth,
+        width: imageElement.naturalHeight,
+      };
 
-  const photosFetched = await Promise.all(photos);
+      return crawledContent;
+    })()
+  );
 
-  const parsedPhotos = photosFetched.filter(item=>item!==null && item.blob.length>=1) as CrawledImageDom[];
+  const photosFetched = await Promise.all(photos);
 
+  const parsedPhotos = photosFetched.filter(
+    (item) => item !== null && item.blob.length >= 1
+  ) as CrawledImageDom[];
 
-  if(parsedPhotos.length >= 1) return parsedPhotos;
+  if (parsedPhotos.length >= 1) return parsedPhotos;
 
-  return []
+  return [];
 }
 
 function listenAddOn() {
   browser.runtime.onMessage.addListener(async () => {
-    const images = await getAllImages();
+    const [images] = await Promise.all([getAllImages()]);
+
     sendImages({
-      action:"crawledContent",
-      message:images
-    })
+      action: "crawledContent",
+      message: {
+        images,
+        urlTab: "",
+      },
+    });
   });
 }
diff --git a/src/customHooks/useAddOn/index.ts b/src/customHooks/useAddOn/index.ts
index 7ca6fdc..caa7fe6 100644
--- a/src/customHooks/useAddOn/index.ts
+++ b/src/customHooks/useAddOn/index.ts
@@ -8,6 +8,7 @@ import { saveAs } from "file-saver";
 
 export default function useAddOn() {
   const [state, setState] = useState<CrawledImageDom[]>([]);
+  const [isTelegramK, setIsTelegramK] = useState<boolean | null>(null);
 
   const global = useContext(ContextGlobal);
 
@@ -20,15 +21,24 @@ export default function useAddOn() {
       const parsedMessage: MessageBrowserActions<"crawledContent"> =
         typeof message === "string" ? JSON.parse(message) : message;
 
-      appendCrawledImage(parsedMessage.message);
+      const images =
+        typeof parsedMessage.message === "string"
+          ? []
+          : parsedMessage.message.images;
+
+      appendCrawledImage(images);
     });
 
     browser.runtime.onMessageExternal.addListener(
       (message: MessageBrowserActions<"crawledContent">) => {
         const parsedMessage: MessageBrowserActions<"crawledContent"> =
           typeof message === "string" ? JSON.parse(message) : message;
+        const images =
+          typeof parsedMessage.message === "string"
+            ? []
+            : parsedMessage.message.images;
 
-        appendCrawledImage(parsedMessage.message);
+        appendCrawledImage(images);
       }
     );
 
@@ -37,9 +47,27 @@ export default function useAddOn() {
       message: "",
     };
 
+    browser.tabs
+      .query({
+        active: true,
+        currentWindow: true,
+      })
+      .then(([tab]) => {
+        const result = checkIsTelegramK(tab.url || "");
+        setIsTelegramK(result);
+        return;
+      });
+
     browser.runtime.sendMessage(JSON.stringify(triggerMessage));
   }, [global]);
 
+  function checkIsTelegramK(url: string): boolean | null {
+    if (url === "") return null;
+
+    const telegramUrlPattern = /^https:\/\/web\.telegram\.org\/k\/.*/;
+    return telegramUrlPattern.test(url);
+  }
+
   function appendCrawledImage(dto: string | CrawledImageDom[]) {
     const crawledContent: CrawledImageDom[] =
       typeof dto === "string" ? JSON.parse(dto) : dto;
@@ -66,5 +94,6 @@ export default function useAddOn() {
   return {
     handleDownloadAll,
     state,
+    isTelegramK,
   };
 }
diff --git a/src/helpers/content_script/types.ts b/src/helpers/content_script/types.ts
index c5d9bf6..7a83131 100644
--- a/src/helpers/content_script/types.ts
+++ b/src/helpers/content_script/types.ts
@@ -1,8 +1,13 @@
 import { CrawledImageDom } from "../../typesContentScript";
 
+export interface CrawlingInformation {
+  images: CrawledImageDom[];
+  urlTab: string;
+}
+
 export type BrowserMessage = "crawledContent" | "popUpOpened";
 
 export interface MessageBrowserActions<T extends BrowserMessage> {
   action: T;
-  message: T extends "crawledContent" ? CrawledImageDom[] | string : string;
+  message: T extends "crawledContent" ? CrawlingInformation | string : string;
 }
diff --git a/src/structure/Gallery/index.tsx b/src/structure/Gallery/index.tsx
index ad1ef81..368af87 100644
--- a/src/structure/Gallery/index.tsx
+++ b/src/structure/Gallery/index.tsx
@@ -9,8 +9,9 @@ export default function Gallery({ items = [] }: PropsGallery) {
   if (items.length <= 0)
     return (
       <div className="p-5">
-        <p>1. Scroll throught your chat</p>
-        <p>2. Once you find media, re-open this popup</p>
+        <p>1. Open a chat</p>
+        <p>2. Navigate throught your chat</p>
+        <p>3. Once you find media, re-open this popup</p>
 
         <hr />
 
diff --git a/src/structure/configuration/index.tsx b/src/structure/configuration/index.tsx
index 2d44ae3..a80dff7 100644
--- a/src/structure/configuration/index.tsx
+++ b/src/structure/configuration/index.tsx
@@ -32,7 +32,8 @@ export default function ConfigurationAddOn() {
           <small>This applys only when you use the "Download All" button</small>
         </div>
 
-        <div className="mb-4">
+        {/* TODO: Uncomment once the funciotnality of content_scripts works as expected */}
+        {/* <div className="mb-4">
           <label htmlFor="displayDownloadOnUi" className="form-label">
             Display download button on chat
           </label>
@@ -50,7 +51,7 @@ export default function ConfigurationAddOn() {
             Displays a download button down each media found on the telegram
             chat
           </small>
-        </div>
+        </div> */}
 
         <div className="mb-4">
           <label htmlFor="avoidThumbnails" className="form-label">