Skip to content

Commit bb358bc

Browse files
patelspratikpatelspratik
committed
fixing ls
1 parent 6237763 commit bb358bc

5 files changed

Lines changed: 201 additions & 36 deletions

File tree

‎pkg/auth/auth.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -348,7 +348,7 @@ func (t Auth) getSavedTokensOrNil() (*entity.AuthTokens, error) {
348348
}
349349
return nil, breverrors.WrapAndTrace(err)
350350
}
351-
if tokens != nil && tokens.AccessToken == "" && tokens.RefreshToken == "" {
351+
if tokens != nil && tokens.AccessToken == "" && tokens.RefreshToken == "" && tokens.APIKey == "" {
352352
return nil, nil
353353
}
354354
return tokens, nil

‎pkg/auth/auth_test.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,28 @@ func TestGetFreshAccessTokenOrNil_APIKeySkipsJWTValidationAndRefresh(t *testing.
107107
assert.False(t, s.didSave)
108108
}
109109

110+
func TestGetFreshAccessTokenOrNil_APIKeyOnlyCredentialReturnsAPIKey(t *testing.T) {
111+
s := MockAuthStore{authTokens: &entity.AuthTokens{
112+
APIKey: "brev_api_test-key",
113+
}}
114+
a := Auth{
115+
&s,
116+
&MockOauth{}, func(_ string) (bool, error) {
117+
t.Fatal("api keys must not be parsed as JWTs")
118+
return false, nil
119+
},
120+
func() (bool, error) {
121+
t.Fatal("api keys must not trigger login")
122+
return false, nil
123+
},
124+
}
125+
126+
res, err := a.GetFreshAccessTokenOrNil()
127+
assert.NoError(t, err)
128+
assert.Equal(t, "brev_api_test-key", res)
129+
assert.False(t, s.didSave)
130+
}
131+
110132
func TestLoginWithAPIKey_SavesTypedCredential(t *testing.T) {
111133
s := MockAuthStore{}
112134
a := Auth{

‎pkg/cmd/ls/ls.go‎

Lines changed: 84 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,14 @@ import (
66
"encoding/json"
77
"fmt"
88
"os"
9+
"strings"
910
"sync"
1011

1112
nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
1213
"connectrpc.com/connect"
1314

1415
"github.com/brevdev/brev-cli/pkg/analytics"
16+
"github.com/brevdev/brev-cli/pkg/auth"
1517
"github.com/brevdev/brev-cli/pkg/externalnode"
1618

1719
"github.com/brevdev/brev-cli/pkg/cmd/cmderrors"
@@ -37,6 +39,7 @@ import (
3739
type LsStore interface {
3840
GetWorkspaces(organizationID string, options *store.GetWorkspacesOptions) ([]entity.Workspace, error)
3941
GetActiveOrganizationOrDefault() (*entity.Organization, error)
42+
GetCachedActiveOrganizationOrNil() (*entity.Organization, error)
4043
GetCurrentUser() (*entity.User, error)
4144
GetUsers(queryParams map[string]string) ([]entity.User, error)
4245
GetWorkspace(workspaceID string) (*entity.Workspace, error)
@@ -75,10 +78,13 @@ with other commands like stop, start, or delete.`,
7578
brev ls orgs --json
7679
`,
7780
PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
81+
if isAPIKeyAuthStore(noLoginLsStore) {
82+
return cmdcontext.InvokeParentPersistentPostRun(cmd, args)
83+
}
7884
if hello.ShouldWeRunOnboardingLSStep(noLoginLsStore) && hello.ShouldWeRunOnboarding(noLoginLsStore) {
7985
// Getting the workspaces should go in the hello.go file but then
8086
// requires passing in stores and that makes it hard to use in other commands
81-
org, err := getOrgForRunLs(loginLsStore, org)
87+
org, err := getOrgForRunLs(loginLsStore, org, false)
8288
if err != nil {
8389
return err
8490
}
@@ -146,9 +152,11 @@ with other commands like stop, start, or delete.`,
146152
// trackLsAnalytics sends analytics event for ls command
147153
func trackLsAnalytics(store LsStore) {
148154
userID := ""
149-
user, err := store.GetCurrentUser()
150-
if err == nil {
151-
userID = user.ID
155+
if !isAPIKeyAuthStore(store) {
156+
user, err := store.GetCurrentUser()
157+
if err == nil {
158+
userID = user.ID
159+
}
152160
}
153161
data := analytics.EventData{
154162
EventName: "Brev ls",
@@ -157,8 +165,30 @@ func trackLsAnalytics(store LsStore) {
157165
_ = analytics.TrackEvent(data)
158166
}
159167

160-
func getOrgForRunLs(lsStore LsStore, orgflag string) (*entity.Organization, error) {
168+
func isAPIKeyAuthStore(lsStore LsStore) bool {
169+
token, err := lsStore.GetAccessToken()
170+
if err != nil {
171+
return false
172+
}
173+
return strings.HasPrefix(strings.TrimSpace(token), auth.BrevAPIKeyPrefix)
174+
}
175+
176+
func getOrgForRunLs(lsStore LsStore, orgflag string, apiKeyAuth bool) (*entity.Organization, error) {
161177
var org *entity.Organization
178+
if apiKeyAuth {
179+
if orgflag != "" {
180+
return nil, breverrors.NewValidationError("api key auth is scoped to the org saved during login; --org is not supported")
181+
}
182+
cachedOrg, err := lsStore.GetCachedActiveOrganizationOrNil()
183+
if err != nil {
184+
return nil, breverrors.WrapAndTrace(err)
185+
}
186+
if cachedOrg == nil || cachedOrg.ID == "" {
187+
return nil, breverrors.NewValidationError("api key auth requires an active org; run brev login --api-key <api-key>::<org-id>")
188+
}
189+
return cachedOrg, nil
190+
}
191+
162192
if orgflag != "" {
163193
var orgs []entity.Organization
164194
orgs, err := lsStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgflag})
@@ -188,12 +218,18 @@ func getOrgForRunLs(lsStore LsStore, orgflag string) (*entity.Organization, erro
188218

189219
func RunLs(t *terminal.Terminal, lsStore LsStore, args []string, orgflag string, showAll bool, jsonOutput bool) error {
190220
ls := NewLs(lsStore, t, jsonOutput)
191-
user, err := lsStore.GetCurrentUser()
192-
if err != nil {
193-
return breverrors.WrapAndTrace(err)
221+
apiKeyAuth := isAPIKeyAuthStore(lsStore)
222+
223+
var user *entity.User
224+
if !apiKeyAuth {
225+
var err error
226+
user, err = lsStore.GetCurrentUser()
227+
if err != nil {
228+
return breverrors.WrapAndTrace(err)
229+
}
194230
}
195231

196-
org, err := getOrgForRunLs(lsStore, orgflag)
232+
org, err := getOrgForRunLs(lsStore, orgflag, apiKeyAuth)
197233
if err != nil {
198234
return breverrors.WrapAndTrace(err)
199235
}
@@ -202,7 +238,7 @@ func RunLs(t *terminal.Terminal, lsStore LsStore, args []string, orgflag string,
202238
}
203239

204240
if len(args) == 1 { //nolint:gocritic // don't want to switch
205-
err = handleLsArg(ls, args[0], user, org, showAll)
241+
err = handleLsArg(ls, args[0], user, org, showAll, apiKeyAuth)
206242
if err != nil {
207243
return breverrors.WrapAndTrace(err)
208244
}
@@ -218,10 +254,13 @@ func RunLs(t *terminal.Terminal, lsStore LsStore, args []string, orgflag string,
218254
return nil
219255
}
220256

221-
func handleLsArg(ls *Ls, arg string, user *entity.User, org *entity.Organization, showAll bool) error {
257+
func handleLsArg(ls *Ls, arg string, user *entity.User, org *entity.Organization, showAll bool, apiKeyAuth bool) error {
222258
// todo refactor this to cmd.register
223259
//nolint:gocritic // idk how to write this as a switch
224260
if util.IsSingularOrPlural(arg, "org") || util.IsSingularOrPlural(arg, "organization") { // handle org, orgs, and organization(s)
261+
if apiKeyAuth {
262+
return breverrors.NewValidationError("api key auth cannot list organizations")
263+
}
225264
err := ls.RunOrgs()
226265
if err != nil {
227266
return breverrors.WrapAndTrace(err)
@@ -232,13 +271,25 @@ func handleLsArg(ls *Ls, arg string, user *entity.User, org *entity.Organization
232271
if err != nil {
233272
return breverrors.WrapAndTrace(err)
234273
}
235-
} else if util.IsSingularOrPlural(arg, "user") && featureflag.IsAdmin(user.GlobalUserType) {
274+
} else if util.IsSingularOrPlural(arg, "user") {
275+
if apiKeyAuth {
276+
return breverrors.NewValidationError("api key auth cannot list users")
277+
}
278+
if !featureflag.IsAdmin(user.GlobalUserType) {
279+
return nil
280+
}
236281
err := ls.RunUser(showAll)
237282
if err != nil {
238283
return breverrors.WrapAndTrace(err)
239284
}
240285
return nil
241-
} else if util.IsSingularOrPlural(arg, "host") && featureflag.IsAdmin(user.GlobalUserType) {
286+
} else if util.IsSingularOrPlural(arg, "host") {
287+
if apiKeyAuth {
288+
return breverrors.NewValidationError("api key auth cannot list hosts")
289+
}
290+
if !featureflag.IsAdmin(user.GlobalUserType) {
291+
return nil
292+
}
242293
err := ls.RunHosts(org)
243294
if err != nil {
244295
return breverrors.WrapAndTrace(err)
@@ -388,6 +439,19 @@ func (ls Ls) ShowUserWorkspaces(org *entity.Organization, otherOrgs []entity.Org
388439
ls.displayWorkspacesAndHelp(org, otherOrgs, userWorkspaces, allWorkspaces, gpuLookup)
389440
}
390441

442+
func (ls Ls) ShowOrgWorkspaces(org *entity.Organization, workspaces []entity.Workspace, gpuLookup map[string]string) {
443+
if len(workspaces) == 0 {
444+
ls.terminal.Vprint(ls.terminal.Yellow("No instances in org %s\n", org.Name))
445+
return
446+
}
447+
ls.terminal.Vprintf("Org %s has %d instances\n", ls.terminal.Yellow(org.Name), len(workspaces))
448+
displayWorkspacesTable(ls.terminal, workspaces, gpuLookup)
449+
450+
fmt.Print("\n")
451+
452+
displayLsResetBreadCrumb(ls.terminal, workspaces)
453+
}
454+
391455
func (ls Ls) displayWorkspacesAndHelp(org *entity.Organization, otherOrgs []entity.Organization, userWorkspaces []entity.Workspace, allWorkspaces []entity.Workspace, gpuLookup map[string]string) {
392456
if len(userWorkspaces) == 0 {
393457
ls.terminal.Vprint(ls.terminal.Yellow("No instances in org %s\n", org.Name))
@@ -489,6 +553,8 @@ func (ls Ls) RunWorkspaces(org *entity.Organization, user *entity.User, showAll
489553
var workspacesToShow []entity.Workspace
490554
if showAll {
491555
workspacesToShow = allWorkspaces
556+
} else if user == nil {
557+
workspacesToShow = allWorkspaces
492558
} else {
493559
workspacesToShow = store.FilterForUserWorkspaces(allWorkspaces, user.ID)
494560
}
@@ -498,6 +564,11 @@ func (ls Ls) RunWorkspaces(org *entity.Organization, user *entity.User, showAll
498564
return ls.outputWorkspacesJSON(workspacesToShow, gpuLookup, nodes)
499565
}
500566

567+
if user == nil {
568+
ls.ShowOrgWorkspaces(org, workspacesToShow, gpuLookup)
569+
return nil
570+
}
571+
501572
// Table output with colors and help text
502573
orgs, err := ls.lsStore.GetOrganizations(nil)
503574
if err != nil {

‎pkg/cmd/ls/ls_test.go‎

Lines changed: 69 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,27 @@ import (
1717
// mockLsStore implements LsStore (including the embedded hello.HelloStore) for
1818
// testing the ls command routing without real API calls.
1919
type mockLsStore struct {
20-
user *entity.User
21-
org *entity.Organization
22-
orgs []entity.Organization
23-
workspaces []entity.Workspace
20+
user *entity.User
21+
org *entity.Organization
22+
orgs []entity.Organization
23+
workspaces []entity.Workspace
24+
accessToken string
25+
currentUserCalls int
26+
getOrganizationsCall int
27+
}
28+
29+
func (m *mockLsStore) GetCurrentUser() (*entity.User, error) {
30+
m.currentUserCalls++
31+
return m.user, nil
32+
}
33+
34+
func (m *mockLsStore) GetAccessToken() (string, error) {
35+
if m.accessToken != "" {
36+
return m.accessToken, nil
37+
}
38+
return "tok", nil
2439
}
2540

26-
func (m *mockLsStore) GetCurrentUser() (*entity.User, error) { return m.user, nil }
27-
func (m *mockLsStore) GetAccessToken() (string, error) { return "tok", nil }
2841
func (m *mockLsStore) GetWorkspace(_ string) (*entity.Workspace, error) {
2942
return nil, nil
3043
}
@@ -33,11 +46,16 @@ func (m *mockLsStore) GetActiveOrganizationOrDefault() (*entity.Organization, er
3346
return m.org, nil
3447
}
3548

49+
func (m *mockLsStore) GetCachedActiveOrganizationOrNil() (*entity.Organization, error) {
50+
return m.org, nil
51+
}
52+
3653
func (m *mockLsStore) GetWorkspaces(_ string, _ *store.GetWorkspacesOptions) ([]entity.Workspace, error) {
3754
return m.workspaces, nil
3855
}
3956

4057
func (m *mockLsStore) GetOrganizations(_ *store.GetOrganizationsOptions) ([]entity.Organization, error) {
58+
m.getOrganizationsCall++
4159
return m.orgs, nil
4260
}
4361

@@ -69,6 +87,49 @@ func newTestStore() *mockLsStore {
6987
}
7088
}
7189

90+
func TestRunLs_APIKeyJSONSkipsUserAndOrgList(t *testing.T) {
91+
s := newTestStore()
92+
s.accessToken = "brev_api_test-key"
93+
s.workspaces = []entity.Workspace{
94+
{
95+
ID: "ws1",
96+
Name: "owned-by-someone",
97+
Status: entity.Running,
98+
CreatedByUserID: "other-user",
99+
},
100+
{
101+
ID: "ws2",
102+
Name: "owned-by-user",
103+
Status: entity.Stopped,
104+
CreatedByUserID: "u1",
105+
},
106+
}
107+
term := terminal.New()
108+
109+
out := captureStdout(t, func() {
110+
err := RunLs(term, s, nil, "", false, true)
111+
if err != nil {
112+
t.Fatalf("RunLs returned error: %v", err)
113+
}
114+
})
115+
116+
var parsed struct {
117+
Workspaces []WorkspaceInfo `json:"workspaces"`
118+
}
119+
if err := json.Unmarshal([]byte(out), &parsed); err != nil {
120+
t.Fatalf("failed to parse JSON output: %v\nraw output: %s", err, out)
121+
}
122+
if len(parsed.Workspaces) != 2 {
123+
t.Fatalf("expected API key ls to show all org workspaces, got %d", len(parsed.Workspaces))
124+
}
125+
if s.currentUserCalls != 0 {
126+
t.Fatalf("expected no GetCurrentUser calls, got %d", s.currentUserCalls)
127+
}
128+
if s.getOrganizationsCall != 0 {
129+
t.Fatalf("expected no GetOrganizations calls, got %d", s.getOrganizationsCall)
130+
}
131+
}
132+
72133
// captureStdout runs fn while capturing stdout and returns the output.
73134
func captureStdout(t *testing.T, fn func()) string {
74135
t.Helper()
@@ -413,15 +474,15 @@ func TestHandleLsArg_Routing(t *testing.T) {
413474
"workspace", "workspaces",
414475
}
415476
for _, arg := range successArgs {
416-
if err := handleLsArg(ls, arg, s.user, s.org, false); err != nil {
477+
if err := handleLsArg(ls, arg, s.user, s.org, false, false); err != nil {
417478
t.Errorf("handleLsArg(%q) returned unexpected error: %v", arg, err)
418479
}
419480
}
420481

421482
// "node"/"nodes" route to RunNodes which calls the gRPC client — verify
422483
// it attempts the path (error expected due to no real client).
423484
for _, arg := range []string{"node", "nodes"} {
424-
_ = handleLsArg(ls, arg, s.user, s.org, false)
485+
_ = handleLsArg(ls, arg, s.user, s.org, false, false)
425486
}
426487
}
427488

0 commit comments

Comments
 (0)