From f08e148b2f88d78896a45da12ac53a4cbe49a022 Mon Sep 17 00:00:00 2001
From: Sudo-Ivan <ivan@quad4.io>
Date: Sat, 27 Dec 2025 03:15:31 -0600
Subject: [PATCH] Add bot blocker functionality and update security middleware

- Introduced a bot blocker initialized with a user agent blocklist.
- Updated the security middleware to utilize the new bot blocker in both main and test files.
- Enhanced error handling for API requests to return a proper 404 response when content is not found.
---
 main.go      | 12 +++++++++++-
 main_test.go |  3 ++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index a39c742..317aa7a 100644
--- a/main.go
+++ b/main.go
@@ -39,6 +39,7 @@ func main() {
 	flag.StringVar(&giteaToken, "t", os.Getenv("GITEA_TOKEN"), "Gitea API Token")
 	flag.StringVar(&giteaServer, "s", "https://git.quad4.io", "Gitea Server URL")
 	flag.StringVar(&configPath, "c", config.DefaultConfigPath, "Path to software.txt (local or remote)")
+	uaBlocklistPath := flag.String("ua-blocklist", "ua-blocklist.txt", "Path to ua-blocklist.txt (optional)")
 	port := flag.String("p", "8080", "Server port")
 	isProd := flag.Bool("prod", os.Getenv("NODE_ENV") == "production", "Run in production mode")
 	updateInterval := flag.Duration("u", 1*time.Hour, "Software update interval")
@@ -49,6 +50,8 @@ func main() {
 	statsService.Start()
 	defer statsService.Stop()
 
+	botBlocker := security.NewBotBlocker(*uaBlocklistPath)
+
 	initialSoftware := config.LoadSoftware(configPath, giteaServer, giteaToken)
 	apiServer := api.NewServer(giteaToken, initialSoftware, statsService)
 	config.StartBackgroundUpdater(configPath, giteaServer, giteaToken, apiServer.SoftwareList.GetLock(), apiServer.SoftwareList.GetDataPtr(), *updateInterval)
@@ -77,7 +80,7 @@ func main() {
 	r.Use(middleware.Logger)
 	r.Use(middleware.Recoverer)
 	r.Use(middleware.Compress(api.CompressionLevel))
-	r.Use(security.SecurityMiddleware(statsService))
+	r.Use(security.SecurityMiddleware(statsService, botBlocker))
 
 	r.Use(httprate.Limit(
 		security.GlobalRateLimit,
@@ -119,6 +122,13 @@ func main() {
 
 		f, err := contentStatic.Open(path)
 		if err != nil {
+			// If it's an API request, return a proper 404
+			if strings.HasPrefix(r.URL.Path, "/api") {
+				http.Error(w, "Not Found", http.StatusNotFound)
+				return
+			}
+
+			// For SPA, serve index.html for unknown frontend routes
 			indexData, err := fs.ReadFile(contentStatic, "index.html")
 			if err != nil {
 				http.Error(w, "Index not found", http.StatusInternalServerError)
diff --git a/main_test.go b/main_test.go
index 26424b6..a73d2bb 100644
--- a/main_test.go
+++ b/main_test.go
@@ -46,11 +46,12 @@ func TestMainHandlers(t *testing.T) {
 
 	giteaServer = mockGitea.URL
 	statsService := stats.NewService("test-hashes.json")
+	botBlocker := security.NewBotBlocker("")
 	initialSoftware := config.LoadSoftware(configPath, giteaServer, "")
 	apiServer := api.NewServer("", initialSoftware, statsService)
 
 	r := chi.NewRouter()
-	r.Use(security.SecurityMiddleware(statsService))
+	r.Use(security.SecurityMiddleware(statsService, botBlocker))
 	r.Get("/api/software", apiServer.APISoftwareHandler)
 	r.Get("/api/stats", statsService.APIStatsHandler)
 	r.Get("/api/download", apiServer.DownloadProxyHandler)