From fbc0738bf6d58235c7c66ec415ea5af3cc791847 Mon Sep 17 00:00:00 2001
From: Joren Schipman <joren.schipman@student.howest.be>
Date: Sat, 4 May 2024 02:07:20 +0200
Subject: [PATCH] Auth

---
 go.mod         |  5 ++++-
 go.sum         |  2 ++
 loothandler.go | 46 ++++++++++++++++++++++++++++++++++++++++------
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/go.mod b/go.mod
index e8bed6e..dac0ff7 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,7 @@ module MalwareServer
 
 go 1.22.2
 
-require github.com/liamg/magic v0.0.1 // indirect
+require (
+	github.com/dgrijalva/jwt-go v3.2.0+incompatible // indirect
+	github.com/liamg/magic v0.0.1 // indirect
+)
diff --git a/go.sum b/go.sum
index 480c891..94afcae 100644
--- a/go.sum
+++ b/go.sum
@@ -1,2 +1,4 @@
+github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM=
+github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/liamg/magic v0.0.1 h1:Ru22ElY+sCh6RvRTWjQzKKCxsEco8hE0co8n1qe7TBM=
 github.com/liamg/magic v0.0.1/go.mod h1:yQkOmZZI52EA+SQ2xyHpVw8fNvTBruF873Y+Vt6S+fk=
diff --git a/loothandler.go b/loothandler.go
index e1fc124..23b03fa 100644
--- a/loothandler.go
+++ b/loothandler.go
@@ -8,19 +8,27 @@ import (
 	"net/http"
 	"path/filepath"
 	"strings"
+	"time"
+
+	"github.com/dgrijalva/jwt-go"
 )
 
 var (
 	password          = "hardcodedpassword"
 	lootPath          = "Loot"
 	sessionCookieName = "auth_session"
-)
+	secretKey         = []byte("key"))
 
 type PageData struct {
 	UIDs  []string
 	Files []string
 }
 
+type Claims struct {
+	Username string `json:"username"`
+	jwt.StandardClaims
+}
+
 func main() {
 	http.HandleFunc("/", logMiddleware(loginHandler))
 	http.HandleFunc("/loot", logMiddleware(lootHandler))
@@ -45,12 +53,28 @@ func loginHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	if r.FormValue("password") == password {
+		expirationTime := time.Now().Add(1 * time.Hour)
+		claims := &Claims{
+			Username: "root",
+			StandardClaims: jwt.StandardClaims{
+				ExpiresAt: expirationTime.Unix(),
+			},
+		}
+
+		token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+		tokenString, err := token.SignedString(secretKey)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
 		http.SetCookie(w, &http.Cookie{
-			Name:   sessionCookieName,
-			Value:  "authenticated",
-			Path:   "/",
-			MaxAge: 3600,
+			Name:    sessionCookieName,
+			Value:   tokenString,
+			Expires: expirationTime,
+			Path:    "/",
 		})
+
 		http.Redirect(w, r, "/loot", http.StatusSeeOther)
 		return
 	}
@@ -100,9 +124,19 @@ func renderTemplate(w http.ResponseWriter, tmpl string, data interface{}) {
 
 func isAuthenticated(r *http.Request) bool {
 	sessionCookie, err := r.Cookie(sessionCookieName)
-	if err != nil || sessionCookie.Value != "authenticated" {
+	if err != nil {
 		return false
 	}
+
+	tokenString := sessionCookie.Value
+	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
+		return secretKey, nil
+	})
+
+	if err != nil || !token.Valid {
+		return false
+	}
+
 	return true
 }