diff --git a/assume-role.go b/assume-role.go
index e0c28a0..74fa220 100644
--- a/assume-role.go
+++ b/assume-role.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rsa"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"libshared"
@@ -28,6 +29,7 @@ type AssumeRoleResponse struct {
 }
 
 var privateKey *rsa.PrivateKey
+var publicKey *rsa.PublicKey
 
 func assumeRole(w http.ResponseWriter, r *http.Request) {
 
@@ -74,36 +76,39 @@ func assumeRole(w http.ResponseWriter, r *http.Request) {
 
 	tokenString := parts[1]
 
-	// Replace with your actual secret or keyfunc
-	secret := []byte("super-secret-key")
+	claims, err := libshared.ValidateJWT(publicKey, tokenString)
 
-	claims := &MyClaims{}
+	if err != nil {
+		log.Println("Invalid token:", err)
 
-	token, err := jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
-		// Validate signing method if needed
-		return secret, nil
-	})
+		msg := "Invalid token"
 
-	if err != nil || !token.Valid {
-		fmt.Println("Invalid token:", err)
-		http.Error(w, "Invalid token", http.StatusUnauthorized)
+		if errors.Is(err, jwt.ErrTokenExpired) {
+			msg = "Token expired"
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusUnauthorized)
+
+		apiresponse := libshared.NewAPIResponse[AssumeRoleResponse]("fail", msg, AssumeRoleResponse{})
+
+		json.NewEncoder(w).Encode(apiresponse)
 		return
 	}
 
-	// Access parsed values
-	log.Println("sub:", claims.Sub)
-	log.Println("purpose:", claims.Purpose)
-	log.Println("account:", claims.Account)
-	// exp and iat come from RegisteredClaims
-	log.Println("exp:", claims.ExpiresAt)
-	log.Println("iat:", claims.IssuedAt)
-
-	if claims.ExpiresAt == nil || time.Now().After(claims.ExpiresAt.Time) {
-		log.Println("Token expired")
-		http.Error(w, "Token expired", http.StatusUnauthorized)
+	claimaccount, ok := claims["account"].(string)
+	if !ok {
+		log.Println("account claim is not a string")
+		// handle error
 		return
 	}
 
+	log.Println("sub:", claims["sub"])
+	log.Println("purpose:", claims["purpose"])
+	log.Println("account:", claimaccount)
+	log.Println("exp:", claims["exp"])
+	log.Println("iat:", claims["iat"])
+
 	// Read the raw request body
 	body, err := io.ReadAll(r.Body)
 
@@ -136,7 +141,7 @@ func assumeRole(w http.ResponseWriter, r *http.Request) {
 
 	checkExisting := libshared.Pool.QueryRow(context.Background(),
 		"SELECT id FROM roles WHERE accountid = $1 AND rolename = $2",
-		claims.Account, role.Rolename)
+		claimaccount, role.Rolename)
 
 	var existingRoleID int64
 	err = checkExisting.Scan(&existingRoleID)
@@ -149,7 +154,7 @@ func assumeRole(w http.ResponseWriter, r *http.Request) {
 
 	fmt.Println("Role ID", existingRoleID)
 
-	roleToken, err := libshared.CreateJWT(privateKey, claims.Account, role.Rolename, "role")
+	roleToken, err := libshared.CreateJWT(privateKey, claimaccount, role.Rolename, "role")
 	if err != nil {
 		log.Println("Error creating JWT:", err)
 		w.WriteHeader(http.StatusInternalServerError)