From aa1057633080c5dfb6fb4a203b42db1340b8a0f8 Mon Sep 17 00:00:00 2001
From: atridadl <me@atri.dad>
Date: Fri, 29 Dec 2023 01:31:36 -0700
Subject: [PATCH] Added allowlist

---
 .env.example   |  4 +++-
 lib/helpers.go | 22 ++++++++++++++++++++++
 main.go        | 19 ++++++-------------
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/.env.example b/.env.example
index 0a6ce53..19d9b5f 100644
--- a/.env.example
+++ b/.env.example
@@ -3,4 +3,6 @@ DISCORD_TOKEN=""
 OPENAI_API_KEY=""
 REPLICATE_API_TOKEN=""
 REDIS_HOST=""
-REDIS_PASSWORD=""
\ No newline at end of file
+REDIS_PASSWORD=""
+# Comma separated
+COOLDOWN_ALLOW_LIST=""
\ No newline at end of file
diff --git a/lib/helpers.go b/lib/helpers.go
index c4a7a9a..10b6e5f 100644
--- a/lib/helpers.go
+++ b/lib/helpers.go
@@ -1,6 +1,9 @@
 package lib
 
 import (
+	"os"
+	"strings"
+
 	"github.com/diamondburned/arikawa/v3/discord"
 )
 
@@ -60,3 +63,22 @@ func GetUserObject(event discord.InteractionEvent) Userish {
 		return directUser{event.User}
 	}
 }
+
+func CooldownHandler(event discord.InteractionEvent) bool {
+	user := GetUserObject(event)
+	allowList := strings.Split(os.Getenv("COOLDOWN_ALLOW_LIST"), ",")
+
+	// Check if the user ID is in the allowList
+	for _, id := range allowList {
+		if id == user.ID().String() {
+			return true
+		}
+	}
+
+	cachedVal := GetCache(user.ID().String() + ":" + "hdpic")
+	if cachedVal != "nil" {
+		return false
+	}
+	SetCache(user.ID().String()+":"+"hdpic", user.ID().String()+":"+"hdpic", 10)
+	return true
+}
diff --git a/main.go b/main.go
index f267dd7..2f94a9e 100644
--- a/main.go
+++ b/main.go
@@ -132,13 +132,11 @@ func (h *handler) cmdPing(ctx context.Context, data cmdroute.CommandData) *api.I
 
 func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData {
 	// Cooldown Logic
-	user := lib.GetUserObject(*data.Event)
+	allowed := lib.CooldownHandler(*data.Event)
 
-	cachedVal := lib.GetCache(user.ID().String() + ":" + "ask")
-	if cachedVal != "nil" {
+	if !allowed {
 		return errorResponse(errors.New("please wait for the cooldown"))
 	}
-	lib.SetCache(user.ID().String()+":"+"ask", user.ID().String()+":"+"ask", 1)
 
 	// Command Logic
 	var options struct {
@@ -180,14 +178,11 @@ func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.In
 
 func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData {
 	// Cooldown Logic
-	user := lib.GetUserObject(*data.Event)
+	allowed := lib.CooldownHandler(*data.Event)
 
-	cachedVal := lib.GetCache(user.ID().String() + ":" + "pic")
-	if cachedVal != "nil" {
+	if !allowed {
 		return errorResponse(errors.New("please wait for the cooldown"))
-
 	}
-	lib.SetCache(user.ID().String()+":"+"pic", user.ID().String()+":"+"pic", 1)
 
 	// Command Logic
 	var options struct {
@@ -259,13 +254,11 @@ func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.In
 
 func (h *handler) cmdHDPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData {
 	// Cooldown Logic
-	user := lib.GetUserObject(*data.Event)
+	allowed := lib.CooldownHandler(*data.Event)
 
-	cachedVal := lib.GetCache(user.ID().String() + ":" + "hdpic")
-	if cachedVal != "nil" {
+	if !allowed {
 		return errorResponse(errors.New("please wait for the cooldown"))
 	}
-	lib.SetCache(user.ID().String()+":"+"hdpic", user.ID().String()+":"+"hdpic", 10)
 
 	// Command Logic
 	var options struct {