package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"github.com/MicahParks/keyfunc"
	"github.com/golang-jwt/jwt/v4"
	"os"
	"strings"
)

func main() {
	region := os.Getenv("COGNITO_REGION")
	poolId := os.Getenv("COGNITO_POOL_ID")

	url := "https://cognito-idp." + region + ".amazonaws.com/" + poolId + "/.well-known/jwks.json"
	tokenFile := os.Args[1]

	tokenBytes, err := os.ReadFile(tokenFile)
	if err != nil {
		writeError(tokenFile, err)
		return
	}
	token := string(tokenBytes)

	jwks, err := keyfunc.Get(url, keyfunc.Options{})
	if err != nil {
		writeError(tokenFile, err)
		return
	}

	parsed, err := jwt.Parse(token, func(t *jwt.Token) (interface{}, error) {
		return jwks.Keyfunc(t)
	})
	if err != nil {
		writeError(tokenFile, err)
		return
	}

	err = parsed.Claims.Valid()
	if err != nil {
		writeError(tokenFile, err)
		return
	}

	parts := strings.Split(parsed.Raw, ".")
	payload := parts[1]
	jsonPayload, err := base64.RawURLEncoding.DecodeString(payload)
	if err != nil {
		writeError(tokenFile, err)
		return
	}

	var payloadData struct {
		Sub string `json:"sub"`
		Exp int    `json:"exp"`
	}
	err = json.NewDecoder(bytes.NewReader(jsonPayload)).Decode(&payloadData)
	if err != nil {
		writeError(tokenFile, err)
		return
	}

	writeResult(tokenFile, Result{
		Success: true,
		Sub:     payloadData.Sub,
		Expiry:  payloadData.Exp,
	})
}

func writeResult(fileName string, result Result) {
	_ = os.Remove(fileName + ".json")

	file, err := os.OpenFile(fileName+".json", os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		panic(err)
	}

	data, err := json.Marshal(result)
	if err != nil {
		panic(err)
	}

	_, err = file.Write(data)
	if err != nil {
		panic(err)
	}
}

func writeError(fileName string, err error) {
	writeResult(fileName, Result{Error: err.Error()})
}

type Result struct {
	Success bool   `json:"success"`
	Error   string `json:"error"`
	Sub     string `json:"sub"`
	Expiry  int    `json:"exp"`
}