From 44a406fd37cf90eb4ef4ee92869d3db2641226ba Mon Sep 17 00:00:00 2001
From: Dimitrij Pinneker <dimitrij.pinneker@capgemini.com>
Date: Fri, 30 Aug 2024 18:08:57 +0000
Subject: [PATCH] Integrated auth

---
 go.mod  |  1 +
 go.sum  |  2 ++
 main.go | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 510d22b..47a33b6 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,7 @@ go 1.21
 
 require (
 	github.com/dutchcoders/go-clamd v0.0.0-20170520113014-b970184f4d9e
+	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/stretchr/testify v1.9.0
 	gopkg.in/gcfg.v1 v1.2.3
 )
diff --git a/go.sum b/go.sum
index 6a6bb1d..ecf68f8 100644
--- a/go.sum
+++ b/go.sum
@@ -3,6 +3,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dutchcoders/go-clamd v0.0.0-20170520113014-b970184f4d9e h1:rcHHSQqzCgvlwP0I/fQ8rQMn/MpHE5gWSLdtpxtP6KQ=
 github.com/dutchcoders/go-clamd v0.0.0-20170520113014-b970184f4d9e/go.mod h1:Byz7q8MSzSPkouskHJhX0er2mZY/m0Vj5bMeMCkkyY4=
+github.com/golang-jwt/jwt/v5 v5.2.1 h1:OuVbFODueb089Lh128TAcimifWaLhJwVflnrgM17wHk=
+github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
diff --git a/main.go b/main.go
index 3733225..b7b9130 100644
--- a/main.go
+++ b/main.go
@@ -23,6 +23,7 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/golang-jwt/jwt/v5"
 	"gopkg.in/gcfg.v1"
 )
 
@@ -120,6 +121,8 @@ var ctx *Ctx
 var configFile string
 var EICAR = []byte(`X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*`)
 
+var secretKey = []byte("secret-key")
+
 func init() {
 	flag.StringVar(&configFile, "config", "", "Configuration file")
 }
@@ -144,7 +147,7 @@ func main() {
 	runtime.GOMAXPROCS(ctx.Config.App.NumThreads)
 
 	startLogging()
-
+	secretKey = []byte(getEnv("CLAMMIT_JWT_SECRET_KEY", "secret-key"))
 	/*
 	 * Construct objects, validate the URLs
 	 */
@@ -164,9 +167,15 @@ func main() {
 	 * Set up the HTTP server
 	 */
 	router := http.NewServeMux()
-
 	router.HandleFunc("/clammit", infoHandler)
-	router.HandleFunc("/clammit/scan", scanHandler)
+	if bytes.Equal(secretKey, []byte("secret-key")) {
+		authenticatedScanHandler := checkAuthentication(http.HandlerFunc(scanHandler))
+		router.HandleFunc("/clammit/scan", func(w http.ResponseWriter, r *http.Request) {
+			authenticatedScanHandler.ServeHTTP(w, r)
+		})
+	} else {
+		router.HandleFunc("/clammit/scan", scanHandler)
+	}
 	router.HandleFunc("/clammit/readyz", readyzHandler)
 
 	if ctx.Config.App.TestPages {
@@ -185,6 +194,23 @@ func main() {
 	}
 }
 
+func checkAuthentication(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cookieName := "jwt"
+		token, err := getTokenFromCookie(w, r, cookieName)
+		if err != nil {
+			log.Fatalf("Not authenticated")
+			return
+		}
+		err2 := verifyToken(token)
+		if err2 != nil {
+			log.Fatalf("Token not valid")
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
 /*
  * Returns the value of an environment variable, or a default value
  */
@@ -231,6 +257,35 @@ func getBoolEnv(key string, fallback bool) bool {
 	return fallback
 }
 
+func getTokenFromCookie(w http.ResponseWriter, r *http.Request, cookieName string) (string, error) {
+	cookie, err := r.Cookie(cookieName)
+	if err != nil {
+		if err == http.ErrNoCookie {
+			w.WriteHeader(http.StatusUnauthorized)
+			return "", err
+		}
+		w.WriteHeader(http.StatusBadRequest)
+		return "", err
+	}
+
+	return cookie.Value, nil
+}
+
+func verifyToken(tokenString string) error {
+	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
+		return secretKey, nil
+	})
+
+	if err != nil {
+		return err
+	}
+
+	if !token.Valid {
+		return fmt.Errorf("invalid token")
+	}
+	return nil
+}
+
 /*
  * Sets the configuration from the file and environment variables
  */