role-manager/assume-role.go

package main

import (
	"context"
	"crypto/rsa"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"libshared"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

/*
 * Returns a temporary token with the permissions of the specified role.
 */

type AssumeRoleRequest struct {
	Rolename string `json:"rolename"`
}

type AssumeRoleResponse struct {
	Token string `json:"token"`
}

var privateKey *rsa.PrivateKey
var publicKey *rsa.PublicKey

func assumeRole(w http.ResponseWriter, r *http.Request) {

	log.Println("Assume Role Request")

	var response APIResponse[string]

	// Only allow POST
	if r.Method != http.MethodPost {

		response.Timestamp = time.Now().Unix()
		response.Status = "error"
		response.Message = "HTTP POST method not allowed"
		//response.Response: "Testing"

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusMethodNotAllowed)
		json.NewEncoder(w).Encode(response)
		return
	}

	// Get JWT from Authorization header
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {

		response.Timestamp = time.Now().Unix()
		response.Status = "error"
		response.Message = "HTTP POST method not allowed"
		//response.Response: "Testing"

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusUnauthorized)
		json.NewEncoder(w).Encode(response)
		return
	}

	parts := strings.SplitN(authHeader, " ", 2)

	if len(parts) != 2 || parts[0] != "Bearer" {
		fmt.Println("Invalid Authorization header format")
		http.Error(w, "Invalid Authorization format", http.StatusUnauthorized)
		return
	}

	tokenString := parts[1]

	claims, err := libshared.ValidateJWT(publicKey, tokenString)

	if err != nil {
		log.Println("Invalid token:", err)

		msg := "Invalid token"

		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
	}

	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)

	if err != nil {
		log.Println("Failed to read body:", err)
		http.Error(w, "Failed to read body", http.StatusBadRequest)
		return
	}

	defer r.Body.Close()

	// Parse JSON into Role struct
	var role AssumeRoleRequest
	if err := json.Unmarshal(body, &role); err != nil {
		log.Println("Invalid JSON:", err)
		http.Error(w, "Invalid JSON", http.StatusBadRequest)
		return
	}

	if role.Rolename == "" {
		log.Println("Rolename is required")
		http.Error(w, "Rolename is required", http.StatusBadRequest)
		return
	}

	// Print the parsed values to terminal
	log.Println("Rolename:", role.Rolename)

	// Prevent duplication of roles with the same name for the same account

	checkExisting := libshared.Pool.QueryRow(context.Background(),
		"SELECT id FROM roles WHERE accountid = $1 AND rolename = $2",
		claimaccount, role.Rolename)

	var existingRoleID int64
	err = checkExisting.Scan(&existingRoleID)

	if err != nil {
		log.Println("Role does not exist", err)
		w.WriteHeader(http.StatusNotFound)
		return
	}

	fmt.Println("Role ID", existingRoleID)

	roleToken, err := libshared.CreateJWT(privateKey, claimaccount, role.Rolename, "role")
	if err != nil {
		log.Println("Error creating JWT:", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	fmt.Println(roleToken)

	apiresponse := libshared.APIResponse[AssumeRoleResponse]{}
	apiresponse.Timestamp = time.Now().Unix()
	apiresponse.Status = "success"
	apiresponse.Message = "Role assumed successfully"
	apiresponse.Content = AssumeRoleResponse{Token: roleToken}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(apiresponse)

	//	w.Write([]byte("JWT parsed successfully"))

}