Add API token integration (#1075)

* add API token integration * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rerun tests * add new workflow * undo linux build * fix from file commands printing extra text * fix script * no secrets on template files * add jwt parser * temporarily hard code cluster * update API token login * no error if no deployments * make sure of API Token * remove comments * fix lint * fix test * fix lint * fix test * fix test * fix test * fix test * fix test * empty commit * update setup * fix lint * fix test * add test * fix lint * add test * fix issues from code review * Update astro-client/client.go Co-authored-by: kushalmalani <[email protected]> * fix tests --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: kushalmalani <[email protected]>
astronomer · Mar 15, 2023 · 2c80d5e · 2c80d5e
1 parent 69e5305
commit 2c80d5e
Show file tree

Hide file tree

Showing 14 changed files with 280 additions and 43 deletions.
diff --git a/astro-client/client.go b/astro-client/client.go
@@ -13,7 +13,7 @@ import (
 )
 
 const (
-	AstronomerConnectionErrMsg = "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions.\n\nDetails"
+	AstronomerConnectionErrMsg = "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions. If you are using an API Key or Token make sure your context is correct.\n\nDetails"
 
 	permissionsErrMsg = "you do not have the appropriate permissions for that"
 )

diff --git a/cloud/deployment/fromfile/fromfile.go b/cloud/deployment/fromfile/fromfile.go
@@ -83,6 +83,7 @@ func CreateOrUpdate(inputFile, action string, client astro.Client, out io.Writer
 	if err != nil {
 		return err
 	}
+
 	existingDeployments, err = client.ListDeployments(c.Organization, workspaceID)
 	if err != nil {
 		return err

diff --git a/cloud/deployment/inspect/inspect.go b/cloud/deployment/inspect/inspect.go
@@ -303,10 +303,14 @@ func getTemplate(formattedDeployment *FormattedDeployment) FormattedDeployment {
 	template := *formattedDeployment
 	template.Deployment.Configuration.Name = ""
 	template.Deployment.Metadata = nil
+	newEnvVars := []EnvironmentVariable{}
 
 	for i := range template.Deployment.EnvVars {
-		// zero out updated at timestamp
-		template.Deployment.EnvVars[i].UpdatedAt = ""
+		if !template.Deployment.EnvVars[i].IsSecret {
+			newEnvVars = append(newEnvVars, template.Deployment.EnvVars[i])
+		}
 	}
+	template.Deployment.EnvVars = newEnvVars
+
 	return template
 }
diff --git a/cloud/deployment/inspect/inspect_test.go b/cloud/deployment/inspect/inspect_test.go
@@ -866,10 +866,8 @@ func TestFormatPrintableDeployment(t *testing.T) {
     environment_variables:
         - is_secret: false
           key: foo
+          updated_at: NOW
           value: bar
-        - is_secret: true
-          key: bar
-          value: baz
     configuration:
         name: ""
         description: description
@@ -1017,12 +1015,8 @@ func TestFormatPrintableDeployment(t *testing.T) {
             {
                 "is_secret": false,
                 "key": "foo",
+                "updated_at": "NOW",
                 "value": "bar"
-            },
-            {
-                "is_secret": true,
-                "key": "bar",
-                "value": "baz"
             }
         ],
         "configuration": {
@@ -1446,9 +1440,17 @@ func TestGetTemplate(t *testing.T) {
 		assert.NoError(t, err)
 		expected.Deployment.Configuration.Name = ""
 		expected.Deployment.Metadata = nil
+		newEnvVars := []EnvironmentVariable{}
 		for i := range expected.Deployment.EnvVars {
-			expected.Deployment.EnvVars[i].UpdatedAt = ""
+			if !expected.Deployment.EnvVars[i].IsSecret {
+				newEnvVars = append(newEnvVars, expected.Deployment.EnvVars[i])
+			}
+		}
+		expected.Deployment.EnvVars = newEnvVars
+		for i := range expected.Deployment.EnvVars {
+			expected.Deployment.EnvVars[i].UpdatedAt = "NOW"
 		}
+
 		actual := getTemplate(&decoded)
 		assert.Equal(t, expected, actual)
 	})
@@ -1473,9 +1475,16 @@ func TestGetTemplate(t *testing.T) {
 		expected.Deployment.Configuration.Name = ""
 		expected.Deployment.Metadata = nil
 		expected.Deployment.EnvVars = nil
+		newEnvVars := []EnvironmentVariable{}
 		for i := range expected.Deployment.EnvVars {
-			expected.Deployment.EnvVars[i].UpdatedAt = ""
+			if !expected.Deployment.EnvVars[i].IsSecret {
+				newEnvVars = append(newEnvVars, expected.Deployment.EnvVars[i])
+			}
 		}
+		for i := range expected.Deployment.EnvVars {
+			expected.Deployment.EnvVars[i].UpdatedAt = "NOW"
+		}
+		expected.Deployment.EnvVars = newEnvVars
 		actual := getTemplate(&decoded)
 		assert.Equal(t, expected, actual)
 	})
@@ -1496,9 +1505,16 @@ func TestGetTemplate(t *testing.T) {
 		expected.Deployment.Configuration.Name = ""
 		expected.Deployment.Metadata = nil
 		expected.Deployment.AlertEmails = nil
+		newEnvVars := []EnvironmentVariable{}
+		for i := range expected.Deployment.EnvVars {
+			if !expected.Deployment.EnvVars[i].IsSecret {
+				newEnvVars = append(newEnvVars, expected.Deployment.EnvVars[i])
+			}
+		}
 		for i := range expected.Deployment.EnvVars {
 			expected.Deployment.EnvVars[i].UpdatedAt = ""
 		}
+		expected.Deployment.EnvVars = newEnvVars
 		actual := getTemplate(&decoded)
 		assert.Equal(t, expected, actual)
 	})

diff --git a/cloud/workspace/workspace_test.go b/cloud/workspace/workspace_test.go
@@ -109,7 +109,7 @@ func TestListError(t *testing.T) {
 
 	buf := new(bytes.Buffer)
 	err := List(astroAPI, buf)
-	assert.EqualError(t, err, "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions.\n\nDetails: Error processing GraphQL request: API error (500): Internal Server Error")
+	assert.EqualError(t, err, "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions. If you are using an API Key or Token make sure your context is correct.\n\nDetails: Error processing GraphQL request: API error (500): Internal Server Error")
 }
 
 func TestGetWorkspaceSelection(t *testing.T) {

diff --git a/cmd/cloud/setup.go b/cmd/cloud/setup.go
@@ -18,20 +18,26 @@ import (
 	"github.com/astronomer/astro-cli/cloud/organization"
 	"github.com/astronomer/astro-cli/context"
 	"github.com/astronomer/astro-cli/pkg/httputil"
+	"github.com/astronomer/astro-cli/pkg/util"
+	"github.com/golang-jwt/jwt/v4"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 )
 
 var (
-	authLogin = auth.Login
-
-	client = httputil.NewHTTPClient()
+	authLogin        = auth.Login
+	defaultDomain    = "astronomer.io"
+	client           = httputil.NewHTTPClient()
+	isDeploymentFile = false
+	parseAPIToken    = util.ParseAPIToken
+	errNotAPIToken   = errors.New("the API token given does not appear to be an Astro API Token")
 )
 
 const (
 	accessTokenExpThreshold = 5 * time.Minute
 	topLvlCmd               = "astro"
+	deploymentCmd           = "deployment"
 )
 
 type TokenResponse struct {
@@ -44,6 +50,18 @@ type TokenResponse struct {
 	ErrorDescription string  `json:"error_description,omitempty"`
 }
 
+type CustomClaims struct {
+	OrgAuthServiceID      string   `json:"org_id"`
+	Scope                 string   `json:"scope"`
+	Permissions           []string `json:"permissions"`
+	Version               string   `json:"version"`
+	IsAstronomerGenerated bool     `json:"isAstronomerGenerated"`
+	RsaKeyID              string   `json:"kid"`
+	APITokenID            string   `json:"apiTokenId"`
+	jwt.RegisteredClaims
+}
+
+//nolint:gocognit
 func Setup(cmd *cobra.Command, args []string, client astro.Client, coreClient astrocore.CoreClient) error {
 	// If the user is trying to login or logout no need to go through auth setup.
 	if cmd.CalledAs() == "login" || cmd.CalledAs() == "logout" {
@@ -79,11 +97,25 @@ func Setup(cmd *cobra.Command, args []string, client astro.Client, coreClient as
 		return nil
 	}
 
+	// if deployment inspect, create, or udpate commands are used
+	deploymentCmds := []string{"inspect", "create", "update"}
+	if util.Contains(deploymentCmds, cmd.CalledAs()) && cmd.Parent().Use == deploymentCmd {
+		isDeploymentFile = true
+	}
+
+	// Check for APITokens before API keys or refresh tokens
+	apiToken, err := checkAPIToken(isDeploymentFile, args)
+	if err != nil {
+		return err
+	}
+	if apiToken {
+		return nil
+	}
+
 	// run auth setup for any command that requires auth
 	apiKey, err := checkAPIKeys(client, coreClient, args)
 	if err != nil {
-		fmt.Println(err)
-		fmt.Println("\nThere was an error using API keys, using regular auth instead")
+		return err
 	}
 	if apiKey {
 		return nil
@@ -207,7 +239,7 @@ func checkAPIKeys(astroClient astro.Client, coreClient astrocore.CoreClient, arg
 	c, err := context.GetCurrentContext() // get current context
 	if err != nil {
 		// set context
-		domain := "astronomer.io"
+		domain := defaultDomain
 		if !context.Exists(domain) {
 			err := context.SetContext(domain)
 			if err != nil {
@@ -308,3 +340,71 @@ func checkAPIKeys(astroClient astro.Client, coreClient astrocore.CoreClient, arg
 	}
 	return true, nil
 }
+
+func checkAPIToken(isDeploymentFile bool, args []string) (bool, error) {
+	// check os variables
+	astroAPIToken := os.Getenv("ASTRO_API_TOKEN")
+	if astroAPIToken == "" {
+		return false, nil
+	}
+	if !isDeploymentFile {
+		fmt.Println("Using an Astro API Token")
+	}
+
+	// get authConfig
+	c, err := context.GetCurrentContext() // get current context
+	if err != nil {
+		// set context
+		domain := defaultDomain
+		if !context.Exists(domain) {
+			err := context.SetContext(domain)
+			if err != nil {
+				return false, err
+			}
+		}
+
+		// Switch context
+		err = context.Switch(domain)
+		if err != nil {
+			return false, err
+		}
+
+		c, err = context.GetContext(domain) // get current context
+		if err != nil {
+			return false, err
+		}
+	}
+
+	err = c.SetContextKey("token", "Bearer "+astroAPIToken)
+	if err != nil {
+		return false, err
+	}
+
+	err = c.SetExpiresIn(time.Now().AddDate(1, 0, 0).Unix())
+	if err != nil {
+		return false, err
+	}
+	// Parse the token to peek at the custom claims
+	claims, err := parseAPIToken(astroAPIToken)
+	if err != nil {
+		return false, err
+	}
+	if len(claims.Permissions) == 0 {
+		return false, errNotAPIToken
+	}
+	workspaceID = strings.Replace(claims.Permissions[1], "workspaceId:", "", 1)
+	orgID := strings.Replace(claims.Permissions[2], "organizationId:", "", 1)
+	orgShortName := strings.Replace(claims.Permissions[3], "orgShortNameId:", "", 1)
+	// If using api keys for virtual runtimes, we dont need to look up for this endpoint
+	if !(len(args) > 0 && strings.HasPrefix(args[0], "vr-")) {
+		err := c.SetContextKey("workspace", workspaceID) // c.Workspace
+		if err != nil {
+			fmt.Println("no workspace set")
+		}
+	}
+	err = c.SetOrganizationContext(orgID, orgShortName)
+	if err != nil {
+		fmt.Println("no organization context set")
+	}
+	return true, nil
+}
diff --git a/cmd/cloud/setup_test.go b/cmd/cloud/setup_test.go
@@ -13,6 +13,7 @@ import (
 	astro_mocks "github.com/astronomer/astro-cli/astro-client/mocks"
 	"github.com/astronomer/astro-cli/context"
 	testUtil "github.com/astronomer/astro-cli/pkg/testing"
+	"github.com/astronomer/astro-cli/pkg/util"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
@@ -111,6 +112,23 @@ func TestSetup(t *testing.T) {
 		assert.NoError(t, err)
 	})
 
+	t.Run("deployment cmd", func(t *testing.T) {
+		testUtil.SetupOSArgsForGinkgo()
+		cmd := &cobra.Command{Use: "inspect"}
+		cmd, err := cmd.ExecuteC()
+		assert.NoError(t, err)
+
+		rootCmd := &cobra.Command{Use: "deployment"}
+		rootCmd.AddCommand(cmd)
+
+		authLogin = func(domain, token string, client astro.Client, coreClient astrocore.CoreClient, out io.Writer, shouldDisplayLoginLink bool) error {
+			return nil
+		}
+
+		err = Setup(cmd, []string{}, nil, nil)
+		assert.NoError(t, err)
+	})
+
 	t.Run("deploy cmd", func(t *testing.T) {
 		testUtil.SetupOSArgsForGinkgo()
 		cmd := &cobra.Command{Use: "deploy"}
@@ -318,3 +336,63 @@ func TestCheckToken(t *testing.T) {
 		assert.Contains(t, err.Error(), "failed to login")
 	})
 }
+
+func TestCheckAPIToken(t *testing.T) {
+	testUtil.InitTestConfig(testUtil.CloudPlatform)
+	t.Run("test context switch", func(t *testing.T) {
+		permissions := []string{
+			"",
+			"workspaceId:workspace-id",
+			"organizationId:org-ID",
+			"orgShortNameId:org-short-name",
+		}
+		mockClaims := util.CustomClaims{
+			Permissions: permissions,
+		}
+
+		authLogin = func(domain, token string, client astro.Client, coreClient astrocore.CoreClient, out io.Writer, shouldDisplayLoginLink bool) error {
+			return nil
+		}
+
+		parseAPIToken = func(astroAPIToken string) (*util.CustomClaims, error) {
+			return &mockClaims, nil
+		}
+
+		t.Setenv("ASTRO_API_TOKEN", "token")
+
+		// Switch context
+		domain := "astronomer-dev.io"
+		err := context.Switch(domain)
+		assert.NoError(t, err)
+
+		// run CheckAPIKeys
+		_, err = checkAPIToken(true, []string{})
+		assert.NoError(t, err)
+	})
+
+	t.Run("bad claims", func(t *testing.T) {
+		permissions := []string{}
+		mockClaims := util.CustomClaims{
+			Permissions: permissions,
+		}
+
+		authLogin = func(domain, token string, client astro.Client, coreClient astrocore.CoreClient, out io.Writer, shouldDisplayLoginLink bool) error {
+			return nil
+		}
+
+		parseAPIToken = func(astroAPIToken string) (*util.CustomClaims, error) {
+			return &mockClaims, nil
+		}
+
+		t.Setenv("ASTRO_API_TOKEN", "token")
+
+		// Switch context
+		domain := "astronomer-dev.io"
+		err := context.Switch(domain)
+		assert.NoError(t, err)
+
+		// run CheckAPIKeys
+		_, err = checkAPIToken(true, []string{})
+		assert.ErrorIs(t, err, errNotAPIToken)
+	})
+}
-Original file line number
+Diff line change
@@ Expand Up / @@ -13,7 +13,7 @@ import ( @@
     )
     const (
-    	AstronomerConnectionErrMsg = "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions.\n\nDetails"
+    	AstronomerConnectionErrMsg = "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions. If you are using an API Key or Token make sure your context is correct.\n\nDetails"
     	permissionsErrMsg = "you do not have the appropriate permissions for that"
     )
@@ Expand Down @@