From fd7f9286b06558b85141b3810bca8f5856f83e8b Mon Sep 17 00:00:00 2001
From: obfuscated-loop <107592590+obfuscated-loop@users.noreply.github.com>
Date: Thu, 14 Mar 2024 00:43:21 +0000
Subject: [PATCH] Improved UA tag search, cleaned code and removed redundancy

---
 middleware.js | 53 ++++++++++++++++++++++++++++++++-------------------
 1 file changed, 33 insertions(+), 20 deletions(-)

diff --git a/middleware.js b/middleware.js
index b8fba5d..359cbc3 100644
--- a/middleware.js
+++ b/middleware.js
@@ -1,29 +1,42 @@
-import { NextRequest, NextResponse, userAgent } from 'next/server';
+import { NextRequest, NextResponse, userAgent } from 'next/server'
 
 const webhook = process.env.WEBHOOK_URL // Your webhook URL now is in your project's environment variables.
 
+// This format allows for multiple different useragent tags to be added per source
+// TODO: Add more sources to this object
+const sourcesToUserAgentTag = {
+  'Discord': ['Discordbot'],
+  'Twitter': ['Twitterbot']
+}
+
 export async function middleware(req) {
-  const ua = userAgent(req)?.ua;
-  const source = ["Mozilla/5.0 (compatible; Discordbot/", "Twitterbot/"].find(u => ua?.startsWith(u))
-  const page = req.url.split("/").slice(-1)[0]
-  
-  await fetch(webhook, {
-    body: JSON.stringify({
-      embeds: [{
-        title: "Triggered view-logger",
-        description: (source ? "Source user-agent: " + ua : "It was loaded an user (or an user on Discord)."),
-        footer: {
-          text: "Requested page: " + page.slice(0, 500),
-        },
-      }],
-    }), headers: { "content-type": "application/json" }, method: "POST"
-  })
+  const userAgentString = userAgent(req)?.ua
+
+  for (const [sourceName, sourceUserAgentTags] of Object.entries(sourcesToUserAgentTag)) {
+    if (sourceUserAgentTags.some(tag => userAgentString.includes(tag))) {
+      var source = sourceName
+      break
+    }
+  }
+
+  // Only send a webhook and image if it's a verified source
+  if (typeof source !== 'undefined') {
+    await fetch(webhook, {
+      body: JSON.stringify({
+        embeds: [{
+          title: 'Triggered view-logger',
+          description: 'A user just viewed your message!',
+          footer: {
+            text: `Source: ${source}`,
+          },
+        }],
+      }), headers: { 'content-type': 'application/json' }, method: 'POST'
+    })
 
-  if (source) {
     // Return the image.
-    return NextResponse.rewrite(new URL("/mini.png", req.url))
+    return NextResponse.rewrite(new URL('/mini.png', req.url))
   } else {
     // Make a message for whoever takes the risk to directly click.
-    return NextResponse.rewrite(new URL("/page.html", req.url));
+    return NextResponse.rewrite(new URL('/page.html', req.url))
   }
-}
+}
\ No newline at end of file