From 3020a9068e89bec1526aeac0f9800f61be215a2f Mon Sep 17 00:00:00 2001
From: "David J. Allen" <allend@lanl.gov>
Date: Thu, 22 Feb 2024 10:57:39 -0700
Subject: [PATCH] Added access token fetching

---
 cmd/config.go       | 46 +++++++++++++++++++++++----------------------
 cmd/login.go        | 18 +++++++++++++++---
 internal/api/api.go | 24 +++++++++++++++++++++++
 3 files changed, 63 insertions(+), 25 deletions(-)

diff --git a/cmd/config.go b/cmd/config.go
index c4d9519..c69b435 100644
--- a/cmd/config.go
+++ b/cmd/config.go
@@ -12,32 +12,34 @@ import (
 )
 
 type Config struct {
-	Host          string   `yaml:"host"`
-	Port          int      `yaml:"port"`
-	RedirectUri   []string `yaml:"redirect-uri"`
-	State         string   `yaml:"state"`
-	ResponseType  string   `yaml:"response-type"`
-	Scope         []string `yaml:"scope"`
-	ClientId      string   `yaml:"client.id"`
-	ClientSecret  string   `yaml:"client.secret"`
-	OIDCHost      string   `yaml:"oidc.host"`
-	OIDCPort      int      `yaml:"oidc.port"`
-	IdentitiesUrl string   `yaml:"identities-url"`
+	Host           string   `yaml:"host"`
+	Port           int      `yaml:"port"`
+	RedirectUri    []string `yaml:"redirect-uri"`
+	State          string   `yaml:"state"`
+	ResponseType   string   `yaml:"response-type"`
+	Scope          []string `yaml:"scope"`
+	ClientId       string   `yaml:"client.id"`
+	ClientSecret   string   `yaml:"client.secret"`
+	OIDCHost       string   `yaml:"oidc.host"`
+	OIDCPort       int      `yaml:"oidc.port"`
+	IdentitiesUrl  string   `yaml:"identities-url"`
+	AccessTokenUrl string   `yaml:"access-token-url"`
 }
 
 func NewConfig() Config {
 	return Config{
-		Host:          "127.0.0.1",
-		Port:          3333,
-		RedirectUri:   []string{""},
-		State:         util.RandomString(20),
-		ResponseType:  "code",
-		Scope:         []string{"openid", "profile", "email"},
-		ClientId:      "",
-		ClientSecret:  "",
-		OIDCHost:      "127.0.0.1",
-		OIDCPort:      80,
-		IdentitiesUrl: "",
+		Host:           "127.0.0.1",
+		Port:           3333,
+		RedirectUri:    []string{""},
+		State:          util.RandomString(20),
+		ResponseType:   "code",
+		Scope:          []string{"openid", "profile", "email"},
+		ClientId:       "",
+		ClientSecret:   "",
+		OIDCHost:       "127.0.0.1",
+		OIDCPort:       80,
+		IdentitiesUrl:  "",
+		AccessTokenUrl: "",
 	}
 }
 
diff --git a/cmd/login.go b/cmd/login.go
index add207e..d0a44fa 100644
--- a/cmd/login.go
+++ b/cmd/login.go
@@ -14,7 +14,8 @@ import (
 )
 
 var (
-	identitiesUrl = ""
+	identitiesUrl  = ""
+	accessTokenUrl = ""
 )
 
 var loginCmd = &cobra.Command{
@@ -29,6 +30,12 @@ var loginCmd = &cobra.Command{
 		oidcProvider := oidc.NewOIDCProvider()
 		oidcProvider.Host = config.OIDCHost
 		oidcProvider.Port = config.OIDCPort
+
+		// check if the client ID is set
+		if config.ClientId == "" {
+			fmt.Printf("client ID must be set\n")
+			os.Exit(1)
+		}
 		var authorizationUrl = util.BuildAuthorizationUrl(
 			oidcProvider.GetAuthorizeUrl(),
 			config.ClientId,
@@ -39,10 +46,10 @@ var loginCmd = &cobra.Command{
 		)
 
 		// print the authorization URL for the user to log in
-		fmt.Printf("Login with identity provider: %s\n", authorizationUrl)
+		fmt.Printf("login with identity provider: %s\n", authorizationUrl)
 
 		// authorize oauth client and listen for callback from provider
-		fmt.Printf("Waiting for response from OIDC provider...\n")
+		fmt.Printf("waiting for response from OIDC provider...\n")
 		code, err := api.WaitForAuthorizationCode(config.Host, config.Port)
 		if errors.Is(err, http.ErrServerClosed) {
 			fmt.Printf("server closed\n")
@@ -67,12 +74,17 @@ var loginCmd = &cobra.Command{
 			api.CreateIdentity(config.IdentitiesUrl, idToken)
 			api.FetchIdentities(config.IdentitiesUrl)
 		}
+
 		// use ID token/user info to get access token from Ory Hydra
+		if config.AccessTokenUrl != "" {
+			api.FetchAccessToken(config.AccessTokenUrl, config.ClientId, idToken)
+		}
 	},
 }
 
 func init() {
 	loginCmd.Flags().StringVar(&config.ClientId, "client.id", config.ClientId, "set the client ID")
+	loginCmd.Flags().StringVar(&config.ClientSecret, "client.secret", config.ClientSecret, "set the client secret")
 	loginCmd.Flags().StringSliceVar(&config.RedirectUri, "redirect-uri", config.RedirectUri, "set the redirect URI")
 	loginCmd.Flags().StringVar(&config.ResponseType, "response-type", config.ResponseType, "set the response-type")
 	loginCmd.Flags().StringSliceVar(&config.Scope, "scope", config.Scope, "set the scopes")
diff --git a/internal/api/api.go b/internal/api/api.go
index 9184010..a770344 100644
--- a/internal/api/api.go
+++ b/internal/api/api.go
@@ -53,6 +53,30 @@ func FetchToken(code string, remoteUrl string, clientId string, clientSecret str
 	return token, nil
 }
 
+func FetchAccessToken(remoteUrl string, clientId string, jwt string) (string, error) {
+	var token string
+	data := url.Values{
+		"grant_type":            {"client_credentials"},
+		"client_id":             {clientId},
+		"client_assertion_type": {"urn:ietf:params:oauth:client-assertion-type:jwt-bearer"},
+		"client_assertion":      {jwt},
+	}
+	res, err := http.PostForm(remoteUrl, data)
+	if err != nil {
+		return "", fmt.Errorf("failed to get token: %s", err)
+	}
+	defer res.Body.Close()
+
+	b, err := io.ReadAll(res.Body)
+	if err != nil {
+		return "", fmt.Errorf("failed to read response body: %v", err)
+	}
+	token = string(b)
+
+	fmt.Printf("%v\n", token)
+	return token, nil
+}
+
 func CreateIdentity(remoteUrl string, idToken string) error {
 	data := []byte(`{
 		"schema_id": "preset://email",