From 35d1d054fa3a570edf08a564cef4b54820ef0ea5 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Wed, 26 Aug 2026 13:29:38 -0700 Subject: [PATCH 1/5] Rename enable-ssh to allow-ssh and add cert-authority support --- go.mod | 4 +- go.sum | 8 +- pkg/cmd/allowssh/allowssh.go | 214 ++++++++++++++++++ .../allowssh_test.go} | 180 +++++++++++++-- pkg/cmd/cmd.go | 6 +- pkg/cmd/deregister/deregister.go | 88 +++++-- pkg/cmd/deregister/deregister_test.go | 163 ++++++++++++- pkg/cmd/disallowssh/disallowssh.go | 99 ++++++++ pkg/cmd/enablessh/enablessh.go | 153 ------------- pkg/cmd/grantssh/grantssh.go | 2 +- pkg/cmd/register/register.go | 12 +- pkg/sshcert/sshcert.go | 55 +++++ pkg/sshcert/sshcert_test.go | 131 +++++++++++ 13 files changed, 912 insertions(+), 203 deletions(-) create mode 100644 pkg/cmd/allowssh/allowssh.go rename pkg/cmd/{enablessh/enablessh_test.go => allowssh/allowssh_test.go} (59%) create mode 100644 pkg/cmd/disallowssh/disallowssh.go delete mode 100644 pkg/cmd/enablessh/enablessh.go diff --git a/go.mod b/go.mod index 5855b1103..83b067310 100644 --- a/go.mod +++ b/go.mod @@ -3,8 +3,8 @@ module github.com/brevdev/brev-cli go 1.25.0 require ( - buf.build/gen/go/brevdev/devplane/connectrpc/go v1.20.0-20260820222245-1cfc91443320.1 - buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260820222245-1cfc91443320.1 + buf.build/gen/go/brevdev/devplane/connectrpc/go v1.20.0-20260827214152-35c65570f2a0.1 + buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260827214152-35c65570f2a0.1 connectrpc.com/connect v1.20.0 github.com/NVIDIA/go-nvml v0.13.0-1 github.com/alessio/shellescape v1.4.1 diff --git a/go.sum b/go.sum index 4af12529f..f8eb670c1 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,7 @@ -buf.build/gen/go/brevdev/devplane/connectrpc/go v1.20.0-20260820222245-1cfc91443320.1 h1:PKIsaGilewnQUSHNUn+Ir4sagWne713vJS3Ys7h9vAY= -buf.build/gen/go/brevdev/devplane/connectrpc/go v1.20.0-20260820222245-1cfc91443320.1/go.mod h1:r4xfuOy9bpAXm13ugDRO+JNmFVlXecGRuKtn1X7os/k= -buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260820222245-1cfc91443320.1 h1:gmAgE9NC+BAovZIs9CNmjgExqM+Gox8AZ6ud3eVMxfA= -buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260820222245-1cfc91443320.1/go.mod h1:N18pnR0HL6srurI7G19FpSEki71wA1u4e2c5zbfeTV8= +buf.build/gen/go/brevdev/devplane/connectrpc/go v1.20.0-20260827214152-35c65570f2a0.1 h1:xzM4gdexDMGgTwdgrlUFHgedW+CSbdXaQzOimV5PbPU= +buf.build/gen/go/brevdev/devplane/connectrpc/go v1.20.0-20260827214152-35c65570f2a0.1/go.mod h1:qMKDH/phd8XN/OWkJlSVhHJ/8P2w1dfE5zUQprRRp8c= +buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260827214152-35c65570f2a0.1 h1:17qqLaEUl7Biv3eyy9owm5yrS/Zi9Zyxc01XSZIzVbU= +buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260827214152-35c65570f2a0.1/go.mod h1:N18pnR0HL6srurI7G19FpSEki71wA1u4e2c5zbfeTV8= buf.build/gen/go/brevdev/protoc-gen-gotag/protocolbuffers/go v1.36.12-20220906235457-8b4922735da5.1 h1:Qk/4GJyWVWvWsfEFeX4T+k7KouZdRUxxUnIUwJ3hmZg= buf.build/gen/go/brevdev/protoc-gen-gotag/protocolbuffers/go v1.36.12-20220906235457-8b4922735da5.1/go.mod h1:SacJAYqnICCQAsBA46cSA/hxhqhxYkiYzseucf6/fhQ= cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= diff --git a/pkg/cmd/allowssh/allowssh.go b/pkg/cmd/allowssh/allowssh.go new file mode 100644 index 000000000..24b886f9c --- /dev/null +++ b/pkg/cmd/allowssh/allowssh.go @@ -0,0 +1,214 @@ +// Package allowssh implements brev allow-ssh. +package allowssh + +import ( + "context" + "fmt" + "os" + "os/exec" + "os/user" + "path/filepath" + "strings" + + nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + "connectrpc.com/connect" + + "github.com/brevdev/brev-cli/pkg/cmd/register" + "github.com/brevdev/brev-cli/pkg/config" + "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/sshcert" + "github.com/brevdev/brev-cli/pkg/terminal" + + "github.com/spf13/cobra" +) + +type AllowSSHStore interface { + GetCurrentUser() (*entity.User, error) + GetAccessToken() (string, error) +} + +type allowSSHDeps struct { + platform externalnode.PlatformChecker + nodeClients externalnode.NodeClientFactory + registrationStore register.RegistrationStore + prompter terminal.Selector + // currentUser resolves the OS user for authorized_keys operations. + currentUser func() (*user.User, error) +} + +func defaultAllowSSHDeps() allowSSHDeps { + return allowSSHDeps{ + platform: register.LinuxPlatform{}, + nodeClients: register.DefaultNodeClientFactory{}, + registrationStore: register.NewFileRegistrationStore(), + prompter: register.TerminalPrompter{}, + currentUser: user.Current, + } +} + +func NewCmdAllowSSH(t *terminal.Terminal, store AllowSSHStore) *cobra.Command { + cmd := &cobra.Command{ + Annotations: map[string]string{"configuration": ""}, + Use: "allow-ssh", + DisableFlagsInUseLine: true, + Short: "Trust the Brev certificate authority on this device for SSH", + Long: "Writes the Brev certificate authority to authorized_keys, allowing this device to be an SSH target for the current Linux user. Users are granted access with 'brev grant-ssh'.", + Example: " brev allow-ssh", + RunE: func(cmd *cobra.Command, args []string) error { + return runAllowSSH(cmd.Context(), t, store, defaultAllowSSHDeps()) + }, + } + + return cmd +} + +func runAllowSSH(ctx context.Context, t *terminal.Terminal, s AllowSSHStore, deps allowSSHDeps) error { + if !deps.platform.IsCompatible() { + return fmt.Errorf("brev allow-ssh is only supported on Linux") + } + + reg, err := deps.registrationStore.Load() + if err != nil { + return fmt.Errorf("failed to read registration file: %w", err) + } + + return allowSSH(ctx, t, deps, s, reg) +} + +func allowSSH( + ctx context.Context, + t *terminal.Terminal, + deps allowSSHDeps, + s AllowSSHStore, + reg *register.DeviceRegistration, +) error { + linuxUser, err := deps.currentUser() + if err != nil { + return fmt.Errorf("failed to determine current Linux user: %w", err) + } + linuxUsername := linuxUser.Username + + checkSSHDaemon(t) + + t.Vprint("") + t.Vprint(t.Green("Allowing SSH on this device")) + t.Vprint("") + t.Vprintf(" Node: %s (%s)\n", reg.DisplayName, reg.ExternalNodeID) + t.Vprintf(" Linux user: %s\n", linuxUsername) + t.Vprint("") + + node, err := fetchRegisteredNode(ctx, deps, s, reg) + if err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + + if node.GetLabels()[sshcert.LabelKeySSHProvider] != sshcert.SSHProviderCertAuth { + return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername) + } + + caPublicKey := node.GetCertificateAuthority() + + if err := installCertAuthority(linuxUser, caPublicKey, reg.ExternalNodeID, linuxUsername); err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + t.Vprint(t.Green(" Certificate authority written to authorized_keys.")) + + t.Vprint("") + t.Vprint(t.Green("SSH allowed on this device. No one has SSH access yet — grant it with: brev grant-ssh")) + return nil +} + +func legacyEnableSSH( + ctx context.Context, + t *terminal.Terminal, + deps allowSSHDeps, + s AllowSSHStore, + reg *register.DeviceRegistration, + node *nodev1.ExternalNode, + linuxUsername string, +) error { + brevUser, err := s.GetCurrentUser() + if err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + + brevPortID, err := register.ResolveSSHAccessPort(ctx, t, deps.prompter, deps.nodeClients, s, reg, node) + if err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + + if err := register.SetupAndRegisterNodeSSHAccess(ctx, t, deps.nodeClients, s, reg, brevUser, linuxUsername, brevPortID); err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + + t.Vprint("") + t.Vprint(t.Green(fmt.Sprintf("SSH access enabled. You can now SSH to this device via: brev shell %s", reg.DisplayName))) + return nil +} + +func installCertAuthority(osUser *user.User, caPublicKey, nodeID, linuxUser string) error { + if caPublicKey == "" { + return fmt.Errorf("certificate authority public key is required") + } + + principal := fmt.Sprintf("brev:v1:vm:%s:login:%s", nodeID, linuxUser) + entry := fmt.Sprintf("cert-authority,principals=\"%s\" %s", principal, strings.TrimSpace(caPublicKey)) + + sshDir := filepath.Join(osUser.HomeDir, ".ssh") + if err := os.MkdirAll(sshDir, 0o700); err != nil { + return fmt.Errorf("creating .ssh directory: %w", err) + } + + authKeysPath := filepath.Join(sshDir, "authorized_keys") + + existing, err := os.ReadFile(authKeysPath) // #nosec G304 + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("reading authorized_keys: %w", err) + } + + // skip if the entry already exists. + for line := range strings.SplitSeq(string(existing), "\n") { + if strings.TrimSpace(line) == entry { + return nil + } + } + + content := string(existing) + if content != "" && !strings.HasSuffix(content, "\n") { + content += "\n" + } + content += entry + "\n" + + if err := os.WriteFile(authKeysPath, []byte(content), 0o600); err != nil { + return fmt.Errorf("writing authorized_keys: %w", err) + } + + return nil +} + +func fetchRegisteredNode( + ctx context.Context, + deps allowSSHDeps, + tokenProvider externalnode.TokenProvider, + reg *register.DeviceRegistration, +) (*nodev1.ExternalNode, error) { + client := deps.nodeClients.NewNodeClient(tokenProvider, config.GlobalConfig.GetBrevPublicAPIURL()) + resp, err := client.GetNode(ctx, connect.NewRequest(&nodev1.GetNodeRequest{ + ExternalNodeId: reg.ExternalNodeID, + })) + if err != nil { + return nil, fmt.Errorf("error retrieving node: %w", err) + } + return resp.Msg.GetExternalNode(), nil +} + +func checkSSHDaemon(t *terminal.Terminal) { + for _, svc := range []string{"ssh", "sshd"} { + out, err := exec.Command("systemctl", "is-active", svc).Output() //nolint:gosec // fixed service names + if err == nil && len(out) > 0 && string(out[:len(out)-1]) == "active" { + return + } + } + t.Vprintf(" %s\n", t.Yellow("Warning: SSH daemon does not appear to be running. SSH access may not work until sshd is started.")) +} diff --git a/pkg/cmd/enablessh/enablessh_test.go b/pkg/cmd/allowssh/allowssh_test.go similarity index 59% rename from pkg/cmd/enablessh/enablessh_test.go rename to pkg/cmd/allowssh/allowssh_test.go index 7df94144d..42ed275c4 100644 --- a/pkg/cmd/enablessh/enablessh_test.go +++ b/pkg/cmd/allowssh/allowssh_test.go @@ -1,7 +1,8 @@ -package enablessh +package allowssh import ( "context" + "fmt" "net/http/httptest" "os" "os/user" @@ -14,16 +15,16 @@ import ( "connectrpc.com/connect" "github.com/brevdev/brev-cli/pkg/cmd/register" + "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/terminal" ) -// tempUser returns a *user.User whose HomeDir points to a temporary directory. func tempUser(t *testing.T) *user.User { t.Helper() return &user.User{HomeDir: t.TempDir()} } -// readAuthorizedKeys is a test helper that reads ~/.ssh/authorized_keys. func readAuthorizedKeys(t *testing.T, u *user.User) string { t.Helper() data, err := os.ReadFile(filepath.Join(u.HomeDir, ".ssh", "authorized_keys")) @@ -224,17 +225,51 @@ func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider return register.NewNodeServiceClient(provider, m.serverURL) } -type mockEnableSSHStore struct { +type mockAllowSSHStore struct { token string } -func (m *mockEnableSSHStore) GetCurrentUser() (interface{}, error) { return nil, nil } -func (m *mockEnableSSHStore) GetAccessToken() (string, error) { return m.token, nil } +func (m *mockAllowSSHStore) GetCurrentUser() (*entity.User, error) { return &entity.User{}, nil } +func (m *mockAllowSSHStore) GetAccessToken() (string, error) { return m.token, nil } + +// mockSelector implements terminal.Selector, returning the first item. +type mockSelector struct{ choice string } + +func (m mockSelector) Select(_ string, items []string) string { + if m.choice != "" { + for _, s := range items { + if s == m.choice { + return s + } + } + } + if len(items) > 0 { + return items[0] + } + return "" +} -// fakeNodeService implements the server side of ExternalNodeService for testing. type fakeNodeService struct { nodev1connect.UnimplementedExternalNodeServiceHandler - getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) + getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) + grantCalls int + openCalls int +} + +func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, _ *connect.Request[nodev1.GrantNodeSSHAccessRequest]) (*connect.Response[nodev1.GrantNodeSSHAccessResponse], error) { + f.grantCalls++ + return connect.NewResponse(&nodev1.GrantNodeSSHAccessResponse{}), nil +} + +func (f *fakeNodeService) OpenPort(_ context.Context, req *connect.Request[nodev1.OpenPortRequest]) (*connect.Response[nodev1.OpenPortResponse], error) { + f.openCalls++ + return connect.NewResponse(&nodev1.OpenPortResponse{ + Port: &nodev1.Port{ + PortId: fmt.Sprintf("port_%d", req.Msg.GetPortNumber()), + Protocol: req.Msg.GetProtocol(), + PortNumber: req.Msg.GetPortNumber(), + }, + }), nil } func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) { @@ -245,14 +280,15 @@ func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1 return connect.NewResponse(resp), nil } -func startFakeServer(t *testing.T, svc *fakeNodeService) (enableSSHDeps, *httptest.Server) { +func startFakeServer(t *testing.T, svc *fakeNodeService) allowSSHDeps { t.Helper() _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) server := httptest.NewServer(handler) t.Cleanup(server.Close) - return enableSSHDeps{ + return allowSSHDeps{ nodeClients: mockNodeClientFactory{serverURL: server.URL}, - }, server + prompter: mockSelector{}, + } } func Test_fetchRegisteredNode(t *testing.T) { @@ -267,8 +303,8 @@ func Test_fetchRegisteredNode(t *testing.T) { }}, nil }, } - deps, _ := startFakeServer(t, svc) - store := &mockEnableSSHStore{token: "tok"} + deps := startFakeServer(t, svc) + store := &mockAllowSSHStore{token: "tok"} reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} node, err := fetchRegisteredNode(context.Background(), deps, store, reg) @@ -279,3 +315,121 @@ func Test_fetchRegisteredNode(t *testing.T) { t.Fatalf("unexpected node: %+v", node) } } + +// --- installCertAuthority --- + +func Test_installCertAuthority(t *testing.T) { + const ( + caKey = "ssh-ed25519 AAAAC3Nz dummyCA" + node = "unode_abc" + luser = "ubuntu" + ) + + t.Run("WritesLine", func(t *testing.T) { + u := tempUser(t) + if err := installCertAuthority(u, caKey, node, luser); err != nil { + t.Fatalf("installCertAuthority: %v", err) + } + want := `cert-authority,principals="brev:v1:vm:unode_abc:login:ubuntu" ssh-ed25519 AAAAC3Nz dummyCA` + if result := readAuthorizedKeys(t, u); !strings.Contains(result, want) { + t.Errorf("expected cert-authority line not found:\n%s", result) + } + }) + + t.Run("Idempotent", func(t *testing.T) { + u := tempUser(t) + for i := 0; i < 2; i++ { + if err := installCertAuthority(u, caKey, node, luser); err != nil { + t.Fatalf("installCertAuthority #%d: %v", i+1, err) + } + } + result := readAuthorizedKeys(t, u) + if n := strings.Count(result, "cert-authority"); n != 1 { + t.Errorf("expected 1 cert-authority line, got %d:\n%s", n, result) + } + }) + + t.Run("PreservesExistingKeys", func(t *testing.T) { + u := tempUser(t) + sshDir := filepath.Join(u.HomeDir, ".ssh") + if err := os.MkdirAll(sshDir, 0o700); err != nil { + t.Fatal(err) + } + original := "ssh-rsa EXISTING user@host\n" + if err := os.WriteFile(filepath.Join(sshDir, "authorized_keys"), []byte(original), 0o600); err != nil { + t.Fatal(err) + } + + if err := installCertAuthority(u, caKey, node, luser); err != nil { + t.Fatalf("installCertAuthority: %v", err) + } + + result := readAuthorizedKeys(t, u) + if !strings.Contains(result, "ssh-rsa EXISTING user@host") { + t.Errorf("existing key was removed:\n%s", result) + } + if !strings.Contains(result, "cert-authority") { + t.Errorf("cert-authority line not written:\n%s", result) + } + }) + + t.Run("EmptyKeyErrors", func(t *testing.T) { + if err := installCertAuthority(tempUser(t), "", node, luser); err == nil { + t.Error("expected error for empty CA key") + } + }) +} + +func Test_allowSSH_LegacyNodeFallsBackToKeys(t *testing.T) { + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: "unode_legacy", + // No sshprovider label — legacy node. + Labels: map[string]string{}, + Ports: []*nodev1.Port{{ + PortId: "port_ssh", + Protocol: nodev1.PortProtocol_PORT_PROTOCOL_TCP, + PortNumber: 22, + }}, + }, + }, nil + }, + } + deps := startFakeServer(t, svc) + // Real username (the legacy path does a system user.Lookup) but temp + // HomeDir so authorized_keys operations never touch the developer's + // real file. + realUser, uerr := user.Current() + if uerr != nil { + t.Fatalf("user.Current: %v", uerr) + } + tempUser := &user.User{Username: realUser.Username, HomeDir: t.TempDir()} + deps.currentUser = func() (*user.User, error) { return tempUser, nil } + + reg := ®ister.DeviceRegistration{ + DisplayName: "legacy-node", + ExternalNodeID: "unode_legacy", + OrgID: "org_1", + } + + term := terminal.New() + if err := allowSSH(context.Background(), term, deps, &mockAllowSSHStore{}, reg); err != nil { + t.Fatalf("allowSSH failed: %v", err) + } + + // Legacy flow must grant SSH access (reflexive grant). + if svc.grantCalls == 0 { + t.Error("expected GrantNodeSSHAccess to be called for legacy node") + } + + // No cert-authority line may be written for a legacy node. + authKeysPath := filepath.Join(tempUser.HomeDir, ".ssh", "authorized_keys") + data, readErr := os.ReadFile(authKeysPath) // #nosec G304 + if readErr == nil { + if strings.Contains(string(data), "brev:v1:vm:unode_legacy") { + t.Errorf("legacy node must not write a cert-authority line:\n%s", string(data)) + } + } +} diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index a845788cd..09ebbff66 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -8,6 +8,7 @@ import ( "github.com/brevdev/brev-cli/pkg/analytics" "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/cmd/agentskill" + "github.com/brevdev/brev-cli/pkg/cmd/allowssh" analyticscmd "github.com/brevdev/brev-cli/pkg/cmd/analytics" "github.com/brevdev/brev-cli/pkg/cmd/background" "github.com/brevdev/brev-cli/pkg/cmd/clipboard" @@ -16,7 +17,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/copy" "github.com/brevdev/brev-cli/pkg/cmd/delete" "github.com/brevdev/brev-cli/pkg/cmd/deregister" - "github.com/brevdev/brev-cli/pkg/cmd/enablessh" + "github.com/brevdev/brev-cli/pkg/cmd/disallowssh" "github.com/brevdev/brev-cli/pkg/cmd/envvars" "github.com/brevdev/brev-cli/pkg/cmd/exec" "github.com/brevdev/brev-cli/pkg/cmd/feedback" @@ -333,7 +334,8 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor cmd.AddCommand(register.NewCmdRegister(t, externalNodeCmdStore)) cmd.AddCommand(deregister.NewCmdDeregister(t, externalNodeCmdStore)) cmd.AddCommand(upgrade.NewCmdUpgrade(t, noLoginCmdStore)) - cmd.AddCommand(enablessh.NewCmdEnableSSH(t, externalNodeCmdStore)) + cmd.AddCommand(allowssh.NewCmdAllowSSH(t, externalNodeCmdStore)) + cmd.AddCommand(disallowssh.NewCmdDisallowSSH(t, externalNodeCmdStore)) cmd.AddCommand(grantssh.NewCmdGrantSSH(t, externalNodeCmdStore)) cmd.AddCommand(revokessh.NewCmdRevokeSSH(t, externalNodeCmdStore)) cmd.AddCommand(runtasks.NewCmdRunTasks(t, noLoginCmdStore)) diff --git a/pkg/cmd/deregister/deregister.go b/pkg/cmd/deregister/deregister.go index eb0d16e13..f05480a7d 100644 --- a/pkg/cmd/deregister/deregister.go +++ b/pkg/cmd/deregister/deregister.go @@ -15,6 +15,7 @@ import ( "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/sshcert" "github.com/brevdev/brev-cli/pkg/sudo" "github.com/brevdev/brev-cli/pkg/terminal" @@ -26,18 +27,27 @@ type DeregisterStore interface { GetAccessToken() (string, error) } -type SSHKeyRemover interface { +type CertAuthorityRemover interface { + RemoveCertAuthority(u *user.User, nodeID, linuxUser string) (bool, error) +} + +// LegacySSHKeyRemover removes Brev-managed per-user SSH keys (legacy nodes). +type LegacySSHKeyRemover interface { RemoveBrevKeys(u *user.User) ([]string, error) } -type brevSSHKeyRemover struct{} +type brevCertAuthorityRemover struct{} + +func (brevCertAuthorityRemover) RemoveCertAuthority(u *user.User, nodeID, linuxUser string) (bool, error) { + removed, err := sshcert.RemoveCertAuthorityLine(u.HomeDir, nodeID, linuxUser) + return removed, breverrors.WrapAndTrace(err) +} + +type legacyKeyRemover struct{} -func (brevSSHKeyRemover) RemoveBrevKeys(u *user.User) ([]string, error) { +func (legacyKeyRemover) RemoveBrevKeys(u *user.User) ([]string, error) { removed, err := register.RemoveBrevAuthorizedKeys(u) - if err != nil { - return nil, fmt.Errorf("removing brev authorized keys: %w", err) - } - return removed, nil + return removed, breverrors.WrapAndTrace(err) } // deregisterDeps bundles the side-effecting dependencies of runDeregister so @@ -50,7 +60,10 @@ type deregisterDeps struct { netbird register.NetBirdManager nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore - sshKeys SSHKeyRemover + sshKeys CertAuthorityRemover + legacyKeys LegacySSHKeyRemover + // currentUser resolves the OS user for authorized_keys operations. + currentUser func() (*user.User, error) } func defaultDeregisterDeps() deregisterDeps { @@ -61,8 +74,10 @@ func defaultDeregisterDeps() deregisterDeps { gater: sudo.Default, netbird: register.Netbird{}, nodeClients: register.DefaultNodeClientFactory{}, + sshKeys: brevCertAuthorityRemover{}, + legacyKeys: legacyKeyRemover{}, registrationStore: register.NewFileRegistrationStore(), - sshKeys: brevSSHKeyRemover{}, + currentUser: user.Current, } } @@ -176,7 +191,7 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, if orgName == "" { orgName = "(unknown)" } - osUser, _ := user.Current() + osUser, _ := deps.currentUser() linuxUser := "(unknown)" if osUser != nil { linuxUser = osUser.Username @@ -197,7 +212,7 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, t.Vprint("") t.Vprint(t.Yellow(" This will:")) t.Vprint(" 1. Remove this node from Brev") - t.Vprint(" 2. Remove Brev SSH keys from this machine (if any)") + t.Vprint(" 2. Remove any SSH data associated with this node") t.Vprint(" 3. Uninstall the Brev tunnel") t.Vprint(" 4. Delete local registration data") t.Vprint("") @@ -213,27 +228,27 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, } } + // a Brev cert-authority line for this node means certauth mode, otherwise legacy per-user keys + certAuth := false + if osUser != nil { + certAuth = sshcert.HasCertAuthorityLine(osUser.HomeDir, reg.ExternalNodeID) + } + t.Vprint(t.Yellow("[Step 1/4] Removing node from Brev...")) if err := removeNodeFromBrev(ctx, t, s, deps, reg); err != nil { return err } t.Vprint("") - t.Vprint(t.Yellow("[Step 2/4] Removing Brev SSH keys...")) + t.Vprint(t.Yellow("[Step 2/4] Removing any SSH data associated with this node...")) if osUser == nil { t.Vprintf(" %s\n", t.Yellow("Skipped: could not determine current user")) } else { - removed, kerr := deps.sshKeys.RemoveBrevKeys(osUser) - switch { - case kerr != nil: - t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to remove Brev SSH keys: %v", kerr))) - case len(removed) > 0: - t.Vprintf("%s Brev SSH keys removed from authorized_keys:\n", t.Green(" ✓")) - for _, key := range removed { - t.Vprintf(" - %s\n", key) - } - default: - t.Vprint(" No Brev SSH keys found in authorized_keys.") + linuxUsername := osUser.Username + if certAuth { + removeCertAuthorityStep(t, deps, osUser, reg.ExternalNodeID, linuxUsername) + } else { + removeLegacyKeysStep(t, deps, osUser) } } t.Vprint("") @@ -260,3 +275,30 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, return nil } + +func removeCertAuthorityStep(t *terminal.Terminal, deps deregisterDeps, osUser *user.User, nodeID, linuxUser string) { + removed, cerr := deps.sshKeys.RemoveCertAuthority(osUser, nodeID, linuxUser) + switch { + case cerr != nil: + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to remove cert-authority: %v", cerr))) + case removed: + t.Vprintf("%s Certificate authority removed from authorized_keys.\n", t.Green(" ✓")) + default: + t.Vprint(" No certificate authority line found in authorized_keys.") + } +} + +func removeLegacyKeysStep(t *terminal.Terminal, deps deregisterDeps, osUser *user.User) { + removed, kerr := deps.legacyKeys.RemoveBrevKeys(osUser) + switch { + case kerr != nil: + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to remove Brev SSH keys: %v", kerr))) + case len(removed) > 0: + t.Vprintf("%s Brev SSH keys removed from authorized_keys:\n", t.Green(" ✓")) + for _, key := range removed { + t.Vprintf(" - %s\n", key) + } + default: + t.Vprint(" No Brev SSH keys found in authorized_keys.") + } +} diff --git a/pkg/cmd/deregister/deregister_test.go b/pkg/cmd/deregister/deregister_test.go index 781d8de48..7c4cdbf69 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -4,7 +4,9 @@ import ( "context" "fmt" "net/http/httptest" + "os" "os/user" + "path/filepath" "strings" "testing" @@ -15,6 +17,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/sshcert" "github.com/brevdev/brev-cli/pkg/sudo" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -39,6 +42,24 @@ type fakeNodeService struct { nodev1connect.UnimplementedExternalNodeServiceHandler removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) listNodesFn func(*nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) + getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) +} + +func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) { + if f.getNodeFn == nil { + // Default: certauth node (matches registration on this branch). + return connect.NewResponse(&nodev1.GetNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: req.Msg.GetExternalNodeId(), + Labels: map[string]string{"sshprovider": "certauth"}, + }, + }), nil + } + resp, err := f.getNodeFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil } func (f *fakeNodeService) RemoveNode(_ context.Context, req *connect.Request[nodev1.RemoveNodeRequest]) (*connect.Response[nodev1.RemoveNodeResponse], error) { @@ -123,12 +144,23 @@ func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider } type mockSSHKeyRemover struct { + called bool + err error + removed bool +} + +func (m *mockSSHKeyRemover) RemoveCertAuthority(_ *user.User, _, _ string) (bool, error) { + m.called = true + return m.removed, m.err +} + +type mockLegacyKeyRemover struct { called bool err error removed []string } -func (m *mockSSHKeyRemover) RemoveBrevKeys(_ *user.User) ([]string, error) { +func (m *mockLegacyKeyRemover) RemoveBrevKeys(_ *user.User) ([]string, error) { m.called = true return m.removed, m.err } @@ -179,6 +211,12 @@ func testDeregisterDeps(t *testing.T, svc *fakeNodeService, regStore register.Re nodeClients: mockNodeClientFactory{serverURL: server.URL}, registrationStore: regStore, sshKeys: &mockSSHKeyRemover{}, + legacyKeys: &mockLegacyKeyRemover{}, + currentUser: func() (*user.User, error) { + // Temp home: tests must never touch the developer's real + // authorized_keys. + return &user.User{HomeDir: t.TempDir(), Username: "testuser"}, nil + }, }, server } @@ -444,6 +482,22 @@ func Test_runDeregister_AlwaysUninstallsNetbird(t *testing.T) { } } +// seedAuthorizedKeys writes a cert-authority line for the given node into a +// fresh temp authorized_keys file and returns the fake user pointing at it. +func seedCertAuthorityUser(t *testing.T, nodeID string) *user.User { + t.Helper() + u := &user.User{HomeDir: t.TempDir(), Username: "testuser"} + sshDir := filepath.Join(u.HomeDir, ".ssh") + if err := os.MkdirAll(sshDir, 0o700); err != nil { + t.Fatal(err) + } + line := fmt.Sprintf("cert-authority,principals=%q ssh-ed25519 TESTCA", sshcert.CertAuthorityPrincipal(nodeID, u.Username)) + if err := os.WriteFile(filepath.Join(sshDir, "authorized_keys"), []byte(line+"\n"), 0o600); err != nil { + t.Fatal(err) + } + return u +} + func Test_runDeregister_RemoveBrevKeysHandling(t *testing.T) { tests := []struct { name string @@ -456,6 +510,10 @@ func Test_runDeregister_RemoveBrevKeysHandling(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + // Seed a local cert-authority line so mode detection picks the + // cert-authority cleanup path this test targets. + tempUser := seedCertAuthorityUser(t, "unode_abc") + regStore := &mockRegistrationStore{reg: registeredReg()} svc := &fakeNodeService{ @@ -464,7 +522,10 @@ func Test_runDeregister_RemoveBrevKeysHandling(t *testing.T) { }, } - err := runDeregisterCase(t, regStore, svc, func(d *deregisterDeps) { d.sshKeys = tt.sshKeys }) + err := runDeregisterCase(t, regStore, svc, func(d *deregisterDeps) { + d.sshKeys = tt.sshKeys + d.currentUser = func() (*user.User, error) { return tempUser, nil } + }) if err != nil { t.Fatalf("runDeregister failed: %v", err) } @@ -484,3 +545,101 @@ func Test_runDeregister_RemoveBrevKeysHandling(t *testing.T) { }) } } + +func Test_runDeregister_LegacyNodeRemovesKeys(t *testing.T) { + regStore := &mockRegistrationStore{reg: registeredReg()} + svc := &fakeNodeService{ + removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + return &nodev1.RemoveNodeResponse{}, nil + }, + getNodeFn: func(req *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: req.GetExternalNodeId(), + // No sshprovider label — legacy node. + Labels: map[string]string{}, + }, + }, nil + }, + } + + certMock := &mockSSHKeyRemover{} + legacyMock := &mockLegacyKeyRemover{removed: []string{"ssh-rsa OLD user@host"}} + + err := runDeregisterCase(t, regStore, svc, func(d *deregisterDeps) { + d.sshKeys = certMock + d.legacyKeys = legacyMock + }) + if err != nil { + t.Fatalf("runDeregister failed: %v", err) + } + + if !legacyMock.called { + t.Error("expected RemoveBrevKeys to be called for legacy node") + } + if certMock.called { + t.Error("expected RemoveCertAuthority NOT to be called for legacy node") + } +} + +func Test_runDeregister_CertAuthNodeRemovesCertAuthority(t *testing.T) { + // Mode detection is local: seed a cert-authority line in a temp + // authorized_keys. + tempUser := seedCertAuthorityUser(t, "unode_abc") + + regStore := &mockRegistrationStore{reg: registeredReg()} + svc := &fakeNodeService{ + removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + return &nodev1.RemoveNodeResponse{}, nil + }, + } + + certMock := &mockSSHKeyRemover{removed: true} + legacyMock := &mockLegacyKeyRemover{} + + err := runDeregisterCase(t, regStore, svc, func(d *deregisterDeps) { + d.sshKeys = certMock + d.legacyKeys = legacyMock + d.currentUser = func() (*user.User, error) { return tempUser, nil } + }) + if err != nil { + t.Fatalf("runDeregister failed: %v", err) + } + + if !certMock.called { + t.Error("expected RemoveCertAuthority to be called for certauth node") + } + if legacyMock.called { + t.Error("expected RemoveBrevKeys NOT to be called for certauth node") + } +} + +func Test_runDeregister_NodeLookupFailure_FallsBackToLocal(t *testing.T) { + regStore := &mockRegistrationStore{reg: registeredReg()} + svc := &fakeNodeService{ + removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + return &nodev1.RemoveNodeResponse{}, nil + }, + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("backend down")) + }, + } + + certMock := &mockSSHKeyRemover{removed: true} + legacyMock := &mockLegacyKeyRemover{removed: []string{"ssh-rsa OLD"}} + + err := runDeregisterCase(t, regStore, svc, func(d *deregisterDeps) { + d.sshKeys = certMock + d.legacyKeys = legacyMock + }) + if err != nil { + t.Fatalf("runDeregister failed: %v", err) + } + + if certMock.called { + t.Error("expected RemoveCertAuthority NOT to be called when no local cert-authority line exists") + } + if !legacyMock.called { + t.Error("expected RemoveBrevKeys to be called on lookup failure") + } +} diff --git a/pkg/cmd/disallowssh/disallowssh.go b/pkg/cmd/disallowssh/disallowssh.go new file mode 100644 index 000000000..fca6d0e0a --- /dev/null +++ b/pkg/cmd/disallowssh/disallowssh.go @@ -0,0 +1,99 @@ +package disallowssh + +import ( + "context" + "fmt" + "os/user" + + "github.com/brevdev/brev-cli/pkg/cmd/register" + "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/sshcert" + "github.com/brevdev/brev-cli/pkg/terminal" + + "github.com/spf13/cobra" +) + +type DisallowSSHStore interface { + GetCurrentUser() (*entity.User, error) + GetAccessToken() (string, error) +} + +type disallowSSHDeps struct { + platform externalnode.PlatformChecker + registrationStore register.RegistrationStore +} + +func defaultDisallowSSHDeps() disallowSSHDeps { + return disallowSSHDeps{ + platform: register.LinuxPlatform{}, + registrationStore: register.NewFileRegistrationStore(), + } +} + +func NewCmdDisallowSSH(t *terminal.Terminal, store DisallowSSHStore) *cobra.Command { + cmd := &cobra.Command{ + Annotations: map[string]string{"configuration": ""}, + Use: "disallow-ssh", + DisableFlagsInUseLine: true, + Short: "Remove Brev SSH access data from this device", + Long: "Removes the Brev certificate authority line and any Brev-managed SSH keys from authorized_keys, revoking SSH access for all users. The node remains registered.", + Example: " brev disallow-ssh", + RunE: func(cmd *cobra.Command, args []string) error { + return runDisallowSSH(cmd.Context(), t, store, defaultDisallowSSHDeps()) + }, + } + + return cmd +} + +func runDisallowSSH(_ context.Context, t *terminal.Terminal, _ DisallowSSHStore, deps disallowSSHDeps) error { + if !deps.platform.IsCompatible() { + return fmt.Errorf("brev disallow-ssh is only supported on Linux") + } + + reg, err := deps.registrationStore.Load() + if err != nil { + return fmt.Errorf("failed to read registration file: %w", err) + } + + linuxUser, err := user.Current() + if err != nil { + return fmt.Errorf("failed to determine current Linux user: %w", err) + } + + t.Vprint("") + t.Vprint(t.Green("Removing SSH certificate authority from this device")) + t.Vprint("") + t.Vprintf(" Node: %s (%s)\n", reg.DisplayName, reg.ExternalNodeID) + t.Vprintf(" Linux user: %s\n", linuxUser.Username) + t.Vprint("") + + removed, err := sshcert.RemoveCertAuthorityLine(linuxUser.HomeDir, reg.ExternalNodeID, linuxUser.Username) + if err != nil { + return fmt.Errorf("disallow SSH failed: %w", err) + } + + if removed { + t.Vprint(t.Green(" Certificate authority removed from authorized_keys.")) + } else { + t.Vprint(t.Yellow(" No certificate authority line found in authorized_keys.")) + } + + // Legacy nodes store per-user keys instead of a cert-authority line. + // Remove them too; both operations are idempotent no-ops when nothing + // matches, so running both covers every node mode. + removedKeys, kerr := register.RemoveBrevAuthorizedKeys(linuxUser) + switch { + case kerr != nil: + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to remove Brev SSH keys: %v", kerr))) + case len(removedKeys) > 0: + t.Vprintf("%s Brev SSH keys removed from authorized_keys:\n", t.Green(" ✓")) + for _, key := range removedKeys { + t.Vprintf(" - %s\n", key) + } + } + + t.Vprint(t.Green("SSH disallowed. Run 'brev allow-ssh' to re-enable.")) + return nil +} diff --git a/pkg/cmd/enablessh/enablessh.go b/pkg/cmd/enablessh/enablessh.go deleted file mode 100644 index 9788b0e6f..000000000 --- a/pkg/cmd/enablessh/enablessh.go +++ /dev/null @@ -1,153 +0,0 @@ -// Package enablessh provides the brev enableSSH command for enabling SSH access -// to a registered external node. -package enablessh - -import ( - "context" - "fmt" - "os/exec" - "os/user" - - nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" - "connectrpc.com/connect" - - "github.com/brevdev/brev-cli/pkg/cmd/register" - "github.com/brevdev/brev-cli/pkg/config" - "github.com/brevdev/brev-cli/pkg/entity" - breverrors "github.com/brevdev/brev-cli/pkg/errors" - "github.com/brevdev/brev-cli/pkg/externalnode" - "github.com/brevdev/brev-cli/pkg/terminal" - - "github.com/spf13/cobra" -) - -// EnableSSHStore defines the store methods needed by the enableSSH command. -type EnableSSHStore interface { - GetCurrentUser() (*entity.User, error) - GetAccessToken() (string, error) -} - -// enableSSHDeps bundles the side-effecting dependencies of runEnableSSH so they -// can be replaced in tests. -type enableSSHDeps struct { - platform externalnode.PlatformChecker - nodeClients externalnode.NodeClientFactory - registrationStore register.RegistrationStore - prompter terminal.Selector -} - -func defaultEnableSSHDeps() enableSSHDeps { - return enableSSHDeps{ - platform: register.LinuxPlatform{}, - nodeClients: register.DefaultNodeClientFactory{}, - registrationStore: register.NewFileRegistrationStore(), - prompter: register.TerminalPrompter{}, - } -} - -func NewCmdEnableSSH(t *terminal.Terminal, store EnableSSHStore) *cobra.Command { - cmd := &cobra.Command{ - Annotations: map[string]string{"configuration": ""}, - Use: "enable-ssh", - DisableFlagsInUseLine: true, - Short: "Enable SSH access to this registered device", - Long: "Enable SSH access to this registered device for the current Brev user.", - Example: " brev enable-ssh", - RunE: func(cmd *cobra.Command, args []string) error { - return runEnableSSH(cmd.Context(), t, store, defaultEnableSSHDeps()) - }, - } - - return cmd -} - -func runEnableSSH(ctx context.Context, t *terminal.Terminal, s EnableSSHStore, deps enableSSHDeps) error { - if !deps.platform.IsCompatible() { - return fmt.Errorf("brev enable-ssh is only supported on Linux") - } - - reg, err := deps.registrationStore.Load() - if err != nil { - return fmt.Errorf("failed to read registration file: %w", err) - } - - brevUser, err := s.GetCurrentUser() - if err != nil { - return breverrors.WrapAndTrace(err) - } - - return enableSSH(ctx, t, deps, s, reg, brevUser) -} - -// enableSSH grants SSH access to the given node for the current Brev user. -// This is the "reflexive grant" — granting yourself SSH access to the device. -func enableSSH( - ctx context.Context, - t *terminal.Terminal, - deps enableSSHDeps, - tokenProvider externalnode.TokenProvider, - reg *register.DeviceRegistration, - brevUser *entity.User, -) error { - linuxUser, err := user.Current() - if err != nil { - return fmt.Errorf("failed to determine current Linux user: %w", err) - } - linuxUsername := linuxUser.Username - - checkSSHDaemon(t) - - t.Vprint("") - t.Vprint(t.Green("Enabling SSH access on this device")) - t.Vprint("") - t.Vprintf(" Node: %s (%s)\n", reg.DisplayName, reg.ExternalNodeID) - t.Vprintf(" Brev user: %s\n", brevUser.ID) - t.Vprintf(" Linux user: %s\n", linuxUsername) - t.Vprint("") - - node, err := fetchRegisteredNode(ctx, deps, tokenProvider, reg) - if err != nil { - return fmt.Errorf("enable SSH failed: %w", err) - } - - brevPortID, err := register.ResolveSSHAccessPort(ctx, t, deps.prompter, deps.nodeClients, tokenProvider, reg, node) - if err != nil { - return fmt.Errorf("enable SSH failed: %w", err) - } - - if err := register.SetupAndRegisterNodeSSHAccess(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, linuxUsername, brevPortID); err != nil { - return fmt.Errorf("enable SSH failed: %w", err) - } - - t.Vprint(t.Green(fmt.Sprintf("SSH access enabled. You can now SSH to this device via: brev shell %s", reg.DisplayName))) - return nil -} - -func fetchRegisteredNode( - ctx context.Context, - deps enableSSHDeps, - tokenProvider externalnode.TokenProvider, - reg *register.DeviceRegistration, -) (*nodev1.ExternalNode, error) { - client := deps.nodeClients.NewNodeClient(tokenProvider, config.GlobalConfig.GetBrevPublicAPIURL()) - resp, err := client.GetNode(ctx, connect.NewRequest(&nodev1.GetNodeRequest{ - ExternalNodeId: reg.ExternalNodeID, - OrganizationId: reg.OrgID, - })) - if err != nil { - return nil, fmt.Errorf("error retrieving node: %w", err) - } - return resp.Msg.GetExternalNode(), nil -} - -// checkSSHDaemon prints a warning if neither "ssh" nor "sshd" systemd services -// appear to be active. It never returns an error — it is best-effort. -func checkSSHDaemon(t *terminal.Terminal) { - for _, svc := range []string{"ssh", "sshd"} { - out, err := exec.Command("systemctl", "is-active", svc).Output() //nolint:gosec // fixed service names - if err == nil && len(out) > 0 && string(out[:len(out)-1]) == "active" { - return - } - } - t.Vprintf(" %s\n", t.Yellow("Warning: SSH daemon does not appear to be running. SSH access may not work until sshd is started.")) -} diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 8d49a274a..7bf7981ab 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -184,7 +184,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt } linuxUserOptions := uniqueLinuxUsersFromNodeSSHAccess(node) if len(linuxUserOptions) == 0 { - return fmt.Errorf("no Linux users on this node yet; run with --linux-user to specify one (e.g. after enable-ssh on the node)") + return fmt.Errorf("no Linux users on this node yet; run with --linux-user to specify one (e.g. after allow-ssh on the node)") } t.Vprint("") linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions) diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index e1a9366e4..e4684d110 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -87,6 +87,8 @@ var ( registerLong = `Register your device with NVIDIA Brev This command registers this machine with Brev and brings up the Brev tunnel. +Registration does not enable SSH; run 'brev allow-ssh' afterwards to allow SSH +on this device, then 'brev grant-ssh' to grant users SSH access. Two modes are supported: • Interactive (default): run 'brev register' with no flags and follow prompts for device name and org. @@ -103,7 +105,10 @@ flow is used.` brev register # Non-interactive (--name and --org required) - brev register --name my-node --org my-org` + brev register --name my-node --org my-org + + # Allow SSH on this device after registering + brev allow-ssh` ) func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { @@ -136,7 +141,7 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { cmd.Flags().StringVarP(&nameFlag, "name", "n", "", "device name (required when using non-interactive mode)") cmd.Flags().IntVarP(&sshPort, "ssh-port", "p", 0, "SSH port (if ssh access is desired)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip all confirmation prompts (assume yes)") - _ = cmd.Flags().MarkDeprecated("ssh-port", "use 'brev enable-ssh' after registration to enable SSH access") + _ = cmd.Flags().MarkDeprecated("ssh-port", "run 'brev allow-ssh' after registration to allow SSH on this device") return cmd } @@ -319,6 +324,7 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore Name: name, DeviceId: deviceID, NodeSpec: toProtoNodeSpec(hwProfile), + Labels: map[string]string{"sshprovider": "certauth"}, })) if err != nil { var connectErr *connect.Error @@ -356,7 +362,7 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore t.Vprintf("%s Registration complete.\n", t.Green(" ✓")) t.Vprint("") - t.Vprintf(" %s\n", t.Green("To enable SSH access to this device, run: brev enable-ssh")) + t.Vprintf(" %s\n", t.Green("To allow SSH on this device, run: brev allow-ssh")) return nil } diff --git a/pkg/sshcert/sshcert.go b/pkg/sshcert/sshcert.go index 84d8342bc..af24e5fa1 100644 --- a/pkg/sshcert/sshcert.go +++ b/pkg/sshcert/sshcert.go @@ -191,3 +191,58 @@ func writeAtomic(fs afero.Fs, path string, data []byte, mode os.FileMode) error } return breverrors.WrapAndTrace(fs.Rename(tmpName, path)) } + +// CertAuthorityPrincipal returns the SSH certificate principal for a node and +// Linux user +func CertAuthorityPrincipal(nodeID, linuxUser string) string { + return fmt.Sprintf("brev:v1:vm:%s:login:%s", nodeID, linuxUser) +} + +// RemoveCertAuthorityLine removes the cert-authority line for the given node +// and Linux user from ~/.ssh/authorized_keys. Returns true if a line was +// removed. Missing file is treated as nothing-to-remove. +func RemoveCertAuthorityLine(homeDir, nodeID, linuxUser string) (bool, error) { + prefix := fmt.Sprintf("cert-authority,principals=%q ", CertAuthorityPrincipal(nodeID, linuxUser)) + + authKeysPath := filepath.Join(homeDir, ".ssh", "authorized_keys") + + existing, err := os.ReadFile(authKeysPath) // #nosec G304 + if err != nil { + if os.IsNotExist(err) { + return false, nil + } + return false, fmt.Errorf("reading authorized_keys: %w", err) + } + + var kept []string + var removed bool + for line := range strings.SplitSeq(string(existing), "\n") { + trimmed := strings.TrimSpace(line) + if strings.HasPrefix(trimmed, prefix) && strings.Contains(trimmed, "cert-authority") { + removed = true + continue + } + kept = append(kept, line) + } + + if !removed { + return false, nil + } + + result := strings.Join(kept, "\n") + if err := os.WriteFile(authKeysPath, []byte(result), 0o600); err != nil { + return false, fmt.Errorf("writing authorized_keys: %w", err) + } + + return true, nil +} + +// HasCertAuthorityLine reports whether authorized_keys contains a Brev +// cert-authority line for the given node (any Linux user). +func HasCertAuthorityLine(homeDir, nodeID string) bool { + data, err := os.ReadFile(filepath.Join(homeDir, ".ssh", "authorized_keys")) // #nosec G304 + if err != nil { + return false + } + return strings.Contains(string(data), "brev:v1:vm:"+nodeID+":") +} diff --git a/pkg/sshcert/sshcert_test.go b/pkg/sshcert/sshcert_test.go index 5e7299b04..f25ea83ea 100644 --- a/pkg/sshcert/sshcert_test.go +++ b/pkg/sshcert/sshcert_test.go @@ -4,6 +4,9 @@ import ( "crypto/ed25519" "crypto/rand" "encoding/pem" + "os" + "os/user" + "path/filepath" "strings" "testing" "time" @@ -234,3 +237,131 @@ func mustGen(t *testing.T) ([]byte, string) { } return priv, pub } + +func testHomeDir(t *testing.T) *user.User { + t.Helper() + return &user.User{HomeDir: t.TempDir()} +} + +func Test_RemoveCertAuthorityLine_RemovesMatchingLine(t *testing.T) { + u := testHomeDir(t) + sshDir := filepath.Join(u.HomeDir, ".ssh") + if err := os.MkdirAll(sshDir, 0o700); err != nil { + t.Fatal(err) + } + + caKey := "ssh-ed25519 AAAAC3Nz dummyCA" + entry := `cert-authority,principals="brev:v1:vm:unode_abc:login:ubuntu" ` + caKey + content := strings.Join([]string{ + "ssh-rsa EXISTING user@host", + entry, + "ssh-ed25519 OTHER admin@server", + "", + }, "\n") + + if err := os.WriteFile(filepath.Join(sshDir, "authorized_keys"), []byte(content), 0o600); err != nil { + t.Fatal(err) + } + + removed, err := RemoveCertAuthorityLine(u.HomeDir, "unode_abc", "ubuntu") + if err != nil { + t.Fatalf("removeCertAuthority: %v", err) + } + if !removed { + t.Fatal("expected line to be removed") + } + + data, err := os.ReadFile(filepath.Join(sshDir, "authorized_keys")) + if err != nil { + t.Fatal(err) + } + + result := string(data) + if strings.Contains(result, caKey) { + t.Errorf("CA key still present:\n%s", result) + } + if strings.Contains(result, "cert-authority") { + t.Errorf("cert-authority line still present:\n%s", result) + } + if !strings.Contains(result, "ssh-rsa EXISTING user@host") { + t.Errorf("non-brev key was removed:\n%s", result) + } +} + +func Test_RemoveCertAuthorityLine_NoopWhenFileDoesNotExist(t *testing.T) { + u := testHomeDir(t) + removed, err := RemoveCertAuthorityLine(u.HomeDir, "unode_abc", "ubuntu") + if err != nil { + t.Fatalf("expected no error for missing file: %v", err) + } + if removed { + t.Error("expected removed=false for missing file") + } +} + +func Test_RemoveCertAuthorityLine_NoopWhenNoMatch(t *testing.T) { + u := testHomeDir(t) + sshDir := filepath.Join(u.HomeDir, ".ssh") + if err := os.MkdirAll(sshDir, 0o700); err != nil { + t.Fatal(err) + } + + original := "ssh-rsa EXISTING user@host\n" + if err := os.WriteFile(filepath.Join(sshDir, "authorized_keys"), []byte(original), 0o600); err != nil { + t.Fatal(err) + } + + removed, err := RemoveCertAuthorityLine(u.HomeDir, "unode_abc", "ubuntu") + if err != nil { + t.Fatalf("removeCertAuthority: %v", err) + } + if removed { + t.Error("expected removed=false when no match") + } + + data, err := os.ReadFile(filepath.Join(sshDir, "authorized_keys")) + if err != nil { + t.Fatal(err) + } + if string(data) != original { + t.Errorf("file was modified when it shouldn't have been") + } +} + +func Test_RemoveCertAuthorityLine_OnlyRemovesMatchingPrincipal(t *testing.T) { + u := testHomeDir(t) + sshDir := filepath.Join(u.HomeDir, ".ssh") + if err := os.MkdirAll(sshDir, 0o700); err != nil { + t.Fatal(err) + } + + otherEntry := `cert-authority,principals="brev:v1:vm:other_node:login:ubuntu" ssh-ed25519 OTHER_CA` + targetEntry := `cert-authority,principals="brev:v1:vm:unode_abc:login:ubuntu" ssh-ed25519 TARGET_CA` + content := strings.Join([]string{ + otherEntry, + targetEntry, + "", + }, "\n") + + if err := os.WriteFile(filepath.Join(sshDir, "authorized_keys"), []byte(content), 0o600); err != nil { + t.Fatal(err) + } + + removed, err := RemoveCertAuthorityLine(u.HomeDir, "unode_abc", "ubuntu") + if err != nil { + t.Fatalf("removeCertAuthority: %v", err) + } + if !removed { + t.Fatal("expected line to be removed") + } + + data, _ := os.ReadFile(filepath.Join(sshDir, "authorized_keys")) + result := string(data) + + if strings.Contains(result, "TARGET_CA") { + t.Errorf("target CA still present:\n%s", result) + } + if !strings.Contains(result, "OTHER_CA") { + t.Errorf("other node's CA was removed:\n%s", result) + } +} From 39d9c3af20ea34dc03a129832c96fe176998d9c7 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Tue, 1 Sep 2026 06:59:59 -0700 Subject: [PATCH 2/5] Allow-ssh also opens port. Grant-ssh allows user to choose themselves. --- pkg/cmd/allowssh/allowssh.go | 55 ++++++++++++++++--- pkg/cmd/grantssh/grantssh.go | 49 ++++++++--------- pkg/cmd/grantssh/grantssh_test.go | 48 +++++++++++++--- pkg/cmd/register/device_registration_store.go | 17 +++--- pkg/cmd/register/providers.go | 5 ++ pkg/cmd/register/register.go | 17 +++--- pkg/terminal/types.go | 5 ++ 7 files changed, 138 insertions(+), 58 deletions(-) diff --git a/pkg/cmd/allowssh/allowssh.go b/pkg/cmd/allowssh/allowssh.go index 24b886f9c..cb78468cd 100644 --- a/pkg/cmd/allowssh/allowssh.go +++ b/pkg/cmd/allowssh/allowssh.go @@ -12,6 +12,7 @@ import ( nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" + breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/config" @@ -98,27 +99,63 @@ func allowSSH( t.Vprintf(" Linux user: %s\n", linuxUsername) t.Vprint("") - node, err := fetchRegisteredNode(ctx, deps, s, reg) - if err != nil { - return fmt.Errorf("allow SSH failed: %w", err) - } - - if node.GetLabels()[sshcert.LabelKeySSHProvider] != sshcert.SSHProviderCertAuth { - return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername) + caPublicKey := reg.CertificateAuthority + if caPublicKey == "" { + node, err := fetchRegisteredNode(ctx, deps, s, reg) + if err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + if node.GetLabels()[sshcert.LabelKeySSHProvider] != sshcert.SSHProviderCertAuth { + return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername) + } + caPublicKey = node.GetCertificateAuthority() } - caPublicKey := node.GetCertificateAuthority() - if err := installCertAuthority(linuxUser, caPublicKey, reg.ExternalNodeID, linuxUsername); err != nil { return fmt.Errorf("allow SSH failed: %w", err) } t.Vprint(t.Green(" Certificate authority written to authorized_keys.")) + if err := ensureSSHPort(ctx, t, deps, s, reg); err != nil { + return fmt.Errorf("allow SSH failed: %w", err) + } + t.Vprint("") t.Vprint(t.Green("SSH allowed on this device. No one has SSH access yet — grant it with: brev grant-ssh")) return nil } +func ensureSSHPort(ctx context.Context, t *terminal.Terminal, deps allowSSHDeps, s AllowSSHStore, reg *register.DeviceRegistration) error { + ports, err := fetchRegisteredNode(ctx, deps, s, reg) + if err != nil { + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Note: could not check existing ports: %v", err))) + } + + if p := findExistingSSHPort(ports); p != nil { + t.Vprintf(" SSH port already allocated (%s).\n", register.FormatPortLabel(p)) + return nil + } + + sshPort, err := register.PromptSSHPort(t) + if err != nil { + return fmt.Errorf("reading SSH port: %w", err) + } + + if _, err := register.OpenSSHPort(ctx, t, deps.nodeClients, s, reg, sshPort); err != nil { + return breverrors.WrapAndTrace(err) + } + return nil +} + +func findExistingSSHPort(node *nodev1.ExternalNode) *nodev1.Port { + for _, p := range node.GetPorts() { + if p.GetPortNumber() == 22 { + return p + } + } + return nil +} + func legacyEnableSSH( ctx context.Context, t *terminal.Terminal, diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 7bf7981ab..0619d4f84 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -38,6 +38,7 @@ type GrantSSHStore interface { // can be replaced in tests. type grantSSHDeps struct { prompter terminal.Selector + inputter terminal.Inputter nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore } @@ -49,6 +50,7 @@ type resolvedMember struct { func defaultGrantSSHDeps() grantSSHDeps { return grantSSHDeps{ prompter: register.TerminalPrompter{}, + inputter: register.TerminalPrompter{}, nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), } @@ -67,10 +69,10 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { Use: "grant-ssh", DisableFlagsInUseLine: true, Short: "Grant SSH access to a node for another org member", - Long: "Grant SSH access to a node for another member of your organization. Interactive: no flags, prompts for org, node, port, and user. Non-interactive: --org, --node, --user, --linux-user, and --port-id required.", + Long: "Grant SSH access to a node for another member of your organization. Interactive: no flags, prompts for org, node, port, user, and Linux user. Non-interactive: --org, --node, --user, --linux-user, and --port-id required.", Example: " brev grant-ssh\n brev grant-ssh --org my-org --node my-node --user user@example.com --linux-user ubuntu --port-id port_abc --approve", RunE: func(cmd *cobra.Command, args []string) error { - interactive := orgFlag == "" && nodeFlag == "" && userFlag == "" + interactive := orgFlag == "" && nodeFlag == "" && userFlag == "" && linuxUser == "" && portIDFlag == "" opts := grantSSHOpts{ interactive: interactive, orgName: orgFlag, @@ -87,7 +89,7 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode)") cmd.Flags().StringVarP(&nodeFlag, "node", "n", "", "node name (required in non-interactive mode)") cmd.Flags().StringVarP(&userFlag, "user", "u", "", "Brev user ID or email to grant (required in non-interactive mode)") - cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (required in non-interactive mode)") + cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (e.g. after allow-ssh on the node)") cmd.Flags().StringVar(&portIDFlag, "port-id", "", "Brev port ID to grant access on (required in non-interactive mode)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip confirmation prompt (assume yes)") @@ -107,11 +109,6 @@ type grantSSHOpts struct { // runGrantSSH runs the grant-ssh flow; the only difference by mode is whether we prompt or use opts. func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opts grantSSHOpts, deps grantSSHDeps) error { //nolint:gocognit,gocyclo,funlen // ok - currentUser, err := s.GetCurrentUser() - if err != nil { - return breverrors.WrapAndTrace(err) - } - if !opts.interactive { if opts.orgName == "" || opts.nodeName == "" || opts.userIDOrEmail == "" { return fmt.Errorf("in non-interactive mode --org, --node, and --user are required") @@ -125,6 +122,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt } var org *entity.Organization + var err error if opts.interactive { allOrgs, listErr := s.ListOrganizations() if listErr != nil { @@ -160,7 +158,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt return breverrors.WrapAndTrace(err) } - orgMembers, err := getOrgMembers(ctx, currentUser, t, s, org.ID) + orgMembers, err := getOrgMembers(ctx, t, s, org.ID) if err != nil { return err } @@ -183,11 +181,20 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt return err } linuxUserOptions := uniqueLinuxUsersFromNodeSSHAccess(node) - if len(linuxUserOptions) == 0 { - return fmt.Errorf("no Linux users on this node yet; run with --linux-user to specify one (e.g. after allow-ssh on the node)") - } t.Vprint("") - linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions) + if len(linuxUserOptions) > 0 { + linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions) + } else { + // No SSH access entries yet (e.g. first grant after allow-ssh): + // the node's Linux users aren't known to Brev, and grant-ssh may + // run off-box, so ask for the target Linux user. + linuxUser = deps.inputter.Input(terminal.PromptContent{ + Label: "Linux username on the node", + ErrorMsg: "linux user is required", + AllowEmpty: false, + }) + linuxUser = strings.TrimSpace(linuxUser) + } } else { selectedUser, err = findUserByIDOrEmail(orgMembers, opts.userIDOrEmail) if err != nil { @@ -289,24 +296,16 @@ func findUserByIDOrEmail(members []resolvedMember, idOrEmail string) (*entity.Us return nil, fmt.Errorf("no org member found matching %q", idOrEmail) } -func getOrgMembers(ctx context.Context, currentUser *entity.User, t *terminal.Terminal, s GrantSSHStore, orgID string) ([]resolvedMember, error) { +func getOrgMembers(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, orgID string) ([]resolvedMember, error) { members, err := s.ListOrganizationMembers(ctx, orgID) if err != nil { return nil, fmt.Errorf("failed to fetch org members: %w", err) } - var otherMembers []*nodev1.OrganizationMember - for _, member := range members { - if member.GetUserId() != currentUser.ID { - otherMembers = append(otherMembers, member) - } - } - - if len(otherMembers) == 0 { - return nil, fmt.Errorf("no other members found in current organization") - } + // The current user is selectable: granting yourself SSH access is the + // normal flow for a fresh node (no SSH access entries yet). var resolved []resolvedMember - for _, m := range otherMembers { + for _, m := range members { memberUser, err := s.GetUserByID(m.GetUserId()) if err != nil { t.Vprintf(" Warning: could not resolve user %s: %v\n", m.GetUserId(), err) diff --git a/pkg/cmd/grantssh/grantssh_test.go b/pkg/cmd/grantssh/grantssh_test.go index 6558f0a39..4910aefba 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -179,11 +179,17 @@ func testGrantSSHDeps(t *testing.T, svc *fakeNodeService, regStore register.Regi } return "" }}, + inputter: mockInputter{value: "testuser"}, nodeClients: mockNodeClientFactory{serverURL: server.URL}, registrationStore: regStore, }, server } +// mockInputter implements terminal.Inputter, returning a fixed value. +type mockInputter struct{ value string } + +func (m mockInputter) Input(_ terminal.PromptContent) string { return m.value } + func Test_runGrantSSH_NotRegistered(t *testing.T) { regStore := &mockRegistrationStore{} // no registration @@ -397,7 +403,7 @@ func Test_runGrantSSH_RPCFailure(t *testing.T) { } } -func Test_runGrantSSH_NoOtherMembers(t *testing.T) { +func Test_runGrantSSH_SelfGrantAllowed(t *testing.T) { regStore := &mockRegistrationStore{ reg: ®ister.DeviceRegistration{ ExternalNodeID: "unode_abc", @@ -406,24 +412,50 @@ func Test_runGrantSSH_NoOtherMembers(t *testing.T) { }, } + self := &entity.User{ID: "user_1", Name: "Prat", Email: "prat@example.com"} store := &mockGrantSSHStore{ - user: &entity.User{ID: "user_1"}, + user: self, org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, token: "tok", members: []*nodev1.OrganizationMember{ - {UserId: "user_1"}, // only current user, no others + {UserId: "user_1"}, // only the current user }, - users: map[string]*entity.User{}, + users: map[string]*entity.User{"user_1": self}, } - svc := &fakeNodeService{} + var gotGrantReq *nodev1.GrantNodeSSHAccessRequest + svc := &fakeNodeService{ + listNodesFn: func(_ *nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) { + return &nodev1.ListNodesResponse{ + Items: []*nodev1.ExternalNode{{ + ExternalNodeId: "unode_abc", + Name: "My Spark", + Ports: []*nodev1.Port{{ + PortId: "port_ssh", + PortNumber: 22, + }}, + }}, + }, nil + }, + grantSSHFn: func(req *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) { + gotGrantReq = req + return &nodev1.GrantNodeSSHAccessResponse{}, nil + }, + } deps, server := testGrantSSHDeps(t, svc, regStore) defer server.Close() term := terminal.New() - opts := grantSSHOpts{interactive: true, skipConfirm: true, linuxUser: "testuser"} + opts := grantSSHOpts{interactive: true, skipConfirm: true, linuxUser: "testuser", portID: "port_ssh"} err := runGrantSSH(context.Background(), term, store, opts, deps) - if err == nil { - t.Fatal("expected error when no other members exist") + if err != nil { + t.Fatalf("runGrantSSH failed: %v", err) + } + + if gotGrantReq == nil { + t.Fatal("expected GrantNodeSSHAccess to be called for self-grant") + } + if gotGrantReq.GetUserId() != "user_1" { + t.Errorf("expected grant for current user, got %s", gotGrantReq.GetUserId()) } } diff --git a/pkg/cmd/register/device_registration_store.go b/pkg/cmd/register/device_registration_store.go index e3938fe5b..abfdbe7f4 100644 --- a/pkg/cmd/register/device_registration_store.go +++ b/pkg/cmd/register/device_registration_store.go @@ -28,14 +28,15 @@ const ( // DeviceRegistration is the persistent identity file for a registered device. // Fields align with the AddNodeResponse from dev-plane. type DeviceRegistration struct { - ExternalNodeID string `json:"external_node_id"` - DisplayName string `json:"display_name"` - OrgID string `json:"org_id"` - OrgName string `json:"org_name"` - DeviceID string `json:"device_id"` - RegisteredAt string `json:"registered_at"` - HardwareProfile HardwareProfile `json:"hardware_profile"` - Status string `json:"status,omitempty"` + ExternalNodeID string `json:"external_node_id"` + DisplayName string `json:"display_name"` + OrgID string `json:"org_id"` + OrgName string `json:"org_name"` + DeviceID string `json:"device_id"` + RegisteredAt string `json:"registered_at"` + HardwareProfile HardwareProfile `json:"hardware_profile"` + Status string `json:"status,omitempty"` + CertificateAuthority string `json:"certificate_authority,omitempty"` } // RegistrationStore defines the contract for persisting device registration data. diff --git a/pkg/cmd/register/providers.go b/pkg/cmd/register/providers.go index a3c27bd52..0c6adf843 100644 --- a/pkg/cmd/register/providers.go +++ b/pkg/cmd/register/providers.go @@ -35,6 +35,11 @@ func (TerminalPrompter) Select(label string, items []string) string { }) } +// Input prompts for free-form text input. +func (TerminalPrompter) Input(pc terminal.PromptContent) string { + return terminal.PromptGetInput(pc) +} + // Netbird handles NetBird installation and uninstallation. type Netbird struct{} diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index e4684d110..1fc4bc3eb 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -338,14 +338,15 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore node := addResp.Msg.GetExternalNode() reg := &DeviceRegistration{ - ExternalNodeID: node.GetExternalNodeId(), - DisplayName: name, - OrgID: org.ID, - OrgName: org.Name, - DeviceID: deviceID, - RegisteredAt: time.Now().UTC().Format(time.RFC3339), - HardwareProfile: *hwProfile, - Status: RegistrationStatusRegistered, + ExternalNodeID: node.GetExternalNodeId(), + DisplayName: name, + OrgID: org.ID, + OrgName: org.Name, + DeviceID: deviceID, + RegisteredAt: time.Now().UTC().Format(time.RFC3339), + HardwareProfile: *hwProfile, + Status: RegistrationStatusRegistered, + CertificateAuthority: node.GetCertificateAuthority(), } t.Vprint("") diff --git a/pkg/terminal/types.go b/pkg/terminal/types.go index e7138531e..d686723aa 100644 --- a/pkg/terminal/types.go +++ b/pkg/terminal/types.go @@ -9,3 +9,8 @@ type Confirmer interface { type Selector interface { Select(label string, items []string) string } + +// Inputter prompts the user for free-form text input. +type Inputter interface { + Input(pc PromptContent) string +} From a6d0bda2e0f09f47eddffdf0547fd2c0fe7cb718 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Tue, 1 Sep 2026 07:55:25 -0700 Subject: [PATCH 3/5] Restore enable-ssh naming --- pkg/cmd/cmd.go | 8 +-- .../disablessh.go} | 26 ++++----- .../allowssh.go => enablessh/enablessh.go} | 56 +++++++++---------- .../enablessh_test.go} | 18 +++--- pkg/cmd/grantssh/grantssh.go | 4 +- pkg/cmd/register/register.go | 8 +-- 6 files changed, 60 insertions(+), 60 deletions(-) rename pkg/cmd/{disallowssh/disallowssh.go => disablessh/disablessh.go} (77%) rename pkg/cmd/{allowssh/allowssh.go => enablessh/enablessh.go} (83%) rename pkg/cmd/{allowssh/allowssh_test.go => enablessh/enablessh_test.go} (95%) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 09ebbff66..4eb80b23d 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -8,7 +8,6 @@ import ( "github.com/brevdev/brev-cli/pkg/analytics" "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/cmd/agentskill" - "github.com/brevdev/brev-cli/pkg/cmd/allowssh" analyticscmd "github.com/brevdev/brev-cli/pkg/cmd/analytics" "github.com/brevdev/brev-cli/pkg/cmd/background" "github.com/brevdev/brev-cli/pkg/cmd/clipboard" @@ -17,7 +16,8 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/copy" "github.com/brevdev/brev-cli/pkg/cmd/delete" "github.com/brevdev/brev-cli/pkg/cmd/deregister" - "github.com/brevdev/brev-cli/pkg/cmd/disallowssh" + "github.com/brevdev/brev-cli/pkg/cmd/disablessh" + "github.com/brevdev/brev-cli/pkg/cmd/enablessh" "github.com/brevdev/brev-cli/pkg/cmd/envvars" "github.com/brevdev/brev-cli/pkg/cmd/exec" "github.com/brevdev/brev-cli/pkg/cmd/feedback" @@ -334,8 +334,8 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor cmd.AddCommand(register.NewCmdRegister(t, externalNodeCmdStore)) cmd.AddCommand(deregister.NewCmdDeregister(t, externalNodeCmdStore)) cmd.AddCommand(upgrade.NewCmdUpgrade(t, noLoginCmdStore)) - cmd.AddCommand(allowssh.NewCmdAllowSSH(t, externalNodeCmdStore)) - cmd.AddCommand(disallowssh.NewCmdDisallowSSH(t, externalNodeCmdStore)) + cmd.AddCommand(enablessh.NewCmdEnableSSH(t, externalNodeCmdStore)) + cmd.AddCommand(disablessh.NewCmdDisableSSH(t, externalNodeCmdStore)) cmd.AddCommand(grantssh.NewCmdGrantSSH(t, externalNodeCmdStore)) cmd.AddCommand(revokessh.NewCmdRevokeSSH(t, externalNodeCmdStore)) cmd.AddCommand(runtasks.NewCmdRunTasks(t, noLoginCmdStore)) diff --git a/pkg/cmd/disallowssh/disallowssh.go b/pkg/cmd/disablessh/disablessh.go similarity index 77% rename from pkg/cmd/disallowssh/disallowssh.go rename to pkg/cmd/disablessh/disablessh.go index fca6d0e0a..e284adac5 100644 --- a/pkg/cmd/disallowssh/disallowssh.go +++ b/pkg/cmd/disablessh/disablessh.go @@ -1,4 +1,4 @@ -package disallowssh +package disablessh import ( "context" @@ -14,42 +14,42 @@ import ( "github.com/spf13/cobra" ) -type DisallowSSHStore interface { +type DisableSSHStore interface { GetCurrentUser() (*entity.User, error) GetAccessToken() (string, error) } -type disallowSSHDeps struct { +type disableSSHDeps struct { platform externalnode.PlatformChecker registrationStore register.RegistrationStore } -func defaultDisallowSSHDeps() disallowSSHDeps { - return disallowSSHDeps{ +func defaultDisableSSHDeps() disableSSHDeps { + return disableSSHDeps{ platform: register.LinuxPlatform{}, registrationStore: register.NewFileRegistrationStore(), } } -func NewCmdDisallowSSH(t *terminal.Terminal, store DisallowSSHStore) *cobra.Command { +func NewCmdDisableSSH(t *terminal.Terminal, store DisableSSHStore) *cobra.Command { cmd := &cobra.Command{ Annotations: map[string]string{"configuration": ""}, - Use: "disallow-ssh", + Use: "disable-ssh", DisableFlagsInUseLine: true, Short: "Remove Brev SSH access data from this device", Long: "Removes the Brev certificate authority line and any Brev-managed SSH keys from authorized_keys, revoking SSH access for all users. The node remains registered.", - Example: " brev disallow-ssh", + Example: " brev disable-ssh", RunE: func(cmd *cobra.Command, args []string) error { - return runDisallowSSH(cmd.Context(), t, store, defaultDisallowSSHDeps()) + return runDisableSSH(cmd.Context(), t, store, defaultDisableSSHDeps()) }, } return cmd } -func runDisallowSSH(_ context.Context, t *terminal.Terminal, _ DisallowSSHStore, deps disallowSSHDeps) error { +func runDisableSSH(_ context.Context, t *terminal.Terminal, _ DisableSSHStore, deps disableSSHDeps) error { if !deps.platform.IsCompatible() { - return fmt.Errorf("brev disallow-ssh is only supported on Linux") + return fmt.Errorf("brev disable-ssh is only supported on Linux") } reg, err := deps.registrationStore.Load() @@ -71,7 +71,7 @@ func runDisallowSSH(_ context.Context, t *terminal.Terminal, _ DisallowSSHStore, removed, err := sshcert.RemoveCertAuthorityLine(linuxUser.HomeDir, reg.ExternalNodeID, linuxUser.Username) if err != nil { - return fmt.Errorf("disallow SSH failed: %w", err) + return fmt.Errorf("disable SSH failed: %w", err) } if removed { @@ -94,6 +94,6 @@ func runDisallowSSH(_ context.Context, t *terminal.Terminal, _ DisallowSSHStore, } } - t.Vprint(t.Green("SSH disallowed. Run 'brev allow-ssh' to re-enable.")) + t.Vprint(t.Green("SSH disabled. Run 'brev enable-ssh' to re-enable.")) return nil } diff --git a/pkg/cmd/allowssh/allowssh.go b/pkg/cmd/enablessh/enablessh.go similarity index 83% rename from pkg/cmd/allowssh/allowssh.go rename to pkg/cmd/enablessh/enablessh.go index cb78468cd..23908d6d1 100644 --- a/pkg/cmd/allowssh/allowssh.go +++ b/pkg/cmd/enablessh/enablessh.go @@ -1,5 +1,5 @@ -// Package allowssh implements brev allow-ssh. -package allowssh +// Package enablessh implements brev enable-ssh. +package enablessh import ( "context" @@ -24,12 +24,12 @@ import ( "github.com/spf13/cobra" ) -type AllowSSHStore interface { +type EnableSSHStore interface { GetCurrentUser() (*entity.User, error) GetAccessToken() (string, error) } -type allowSSHDeps struct { +type enableSSHDeps struct { platform externalnode.PlatformChecker nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore @@ -38,8 +38,8 @@ type allowSSHDeps struct { currentUser func() (*user.User, error) } -func defaultAllowSSHDeps() allowSSHDeps { - return allowSSHDeps{ +func defaultEnableSSHDeps() enableSSHDeps { + return enableSSHDeps{ platform: register.LinuxPlatform{}, nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), @@ -48,25 +48,25 @@ func defaultAllowSSHDeps() allowSSHDeps { } } -func NewCmdAllowSSH(t *terminal.Terminal, store AllowSSHStore) *cobra.Command { +func NewCmdEnableSSH(t *terminal.Terminal, store EnableSSHStore) *cobra.Command { cmd := &cobra.Command{ Annotations: map[string]string{"configuration": ""}, - Use: "allow-ssh", + Use: "enable-ssh", DisableFlagsInUseLine: true, Short: "Trust the Brev certificate authority on this device for SSH", Long: "Writes the Brev certificate authority to authorized_keys, allowing this device to be an SSH target for the current Linux user. Users are granted access with 'brev grant-ssh'.", - Example: " brev allow-ssh", + Example: " brev enable-ssh", RunE: func(cmd *cobra.Command, args []string) error { - return runAllowSSH(cmd.Context(), t, store, defaultAllowSSHDeps()) + return runEnableSSH(cmd.Context(), t, store, defaultEnableSSHDeps()) }, } return cmd } -func runAllowSSH(ctx context.Context, t *terminal.Terminal, s AllowSSHStore, deps allowSSHDeps) error { +func runEnableSSH(ctx context.Context, t *terminal.Terminal, s EnableSSHStore, deps enableSSHDeps) error { if !deps.platform.IsCompatible() { - return fmt.Errorf("brev allow-ssh is only supported on Linux") + return fmt.Errorf("brev enable-ssh is only supported on Linux") } reg, err := deps.registrationStore.Load() @@ -74,14 +74,14 @@ func runAllowSSH(ctx context.Context, t *terminal.Terminal, s AllowSSHStore, dep return fmt.Errorf("failed to read registration file: %w", err) } - return allowSSH(ctx, t, deps, s, reg) + return enableSSH(ctx, t, deps, s, reg) } -func allowSSH( +func enableSSH( ctx context.Context, t *terminal.Terminal, - deps allowSSHDeps, - s AllowSSHStore, + deps enableSSHDeps, + s EnableSSHStore, reg *register.DeviceRegistration, ) error { linuxUser, err := deps.currentUser() @@ -93,7 +93,7 @@ func allowSSH( checkSSHDaemon(t) t.Vprint("") - t.Vprint(t.Green("Allowing SSH on this device")) + t.Vprint(t.Green("Enabling SSH on this device")) t.Vprint("") t.Vprintf(" Node: %s (%s)\n", reg.DisplayName, reg.ExternalNodeID) t.Vprintf(" Linux user: %s\n", linuxUsername) @@ -103,7 +103,7 @@ func allowSSH( if caPublicKey == "" { node, err := fetchRegisteredNode(ctx, deps, s, reg) if err != nil { - return fmt.Errorf("allow SSH failed: %w", err) + return fmt.Errorf("enable SSH failed: %w", err) } if node.GetLabels()[sshcert.LabelKeySSHProvider] != sshcert.SSHProviderCertAuth { return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername) @@ -112,20 +112,20 @@ func allowSSH( } if err := installCertAuthority(linuxUser, caPublicKey, reg.ExternalNodeID, linuxUsername); err != nil { - return fmt.Errorf("allow SSH failed: %w", err) + return fmt.Errorf("enable SSH failed: %w", err) } t.Vprint(t.Green(" Certificate authority written to authorized_keys.")) if err := ensureSSHPort(ctx, t, deps, s, reg); err != nil { - return fmt.Errorf("allow SSH failed: %w", err) + return fmt.Errorf("enable SSH failed: %w", err) } t.Vprint("") - t.Vprint(t.Green("SSH allowed on this device. No one has SSH access yet — grant it with: brev grant-ssh")) + t.Vprint(t.Green("SSH enabled on this device. No one has SSH access yet — grant it with: brev grant-ssh")) return nil } -func ensureSSHPort(ctx context.Context, t *terminal.Terminal, deps allowSSHDeps, s AllowSSHStore, reg *register.DeviceRegistration) error { +func ensureSSHPort(ctx context.Context, t *terminal.Terminal, deps enableSSHDeps, s EnableSSHStore, reg *register.DeviceRegistration) error { ports, err := fetchRegisteredNode(ctx, deps, s, reg) if err != nil { t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Note: could not check existing ports: %v", err))) @@ -159,24 +159,24 @@ func findExistingSSHPort(node *nodev1.ExternalNode) *nodev1.Port { func legacyEnableSSH( ctx context.Context, t *terminal.Terminal, - deps allowSSHDeps, - s AllowSSHStore, + deps enableSSHDeps, + s EnableSSHStore, reg *register.DeviceRegistration, node *nodev1.ExternalNode, linuxUsername string, ) error { brevUser, err := s.GetCurrentUser() if err != nil { - return fmt.Errorf("allow SSH failed: %w", err) + return fmt.Errorf("enable SSH failed: %w", err) } brevPortID, err := register.ResolveSSHAccessPort(ctx, t, deps.prompter, deps.nodeClients, s, reg, node) if err != nil { - return fmt.Errorf("allow SSH failed: %w", err) + return fmt.Errorf("enable SSH failed: %w", err) } if err := register.SetupAndRegisterNodeSSHAccess(ctx, t, deps.nodeClients, s, reg, brevUser, linuxUsername, brevPortID); err != nil { - return fmt.Errorf("allow SSH failed: %w", err) + return fmt.Errorf("enable SSH failed: %w", err) } t.Vprint("") @@ -226,7 +226,7 @@ func installCertAuthority(osUser *user.User, caPublicKey, nodeID, linuxUser stri func fetchRegisteredNode( ctx context.Context, - deps allowSSHDeps, + deps enableSSHDeps, tokenProvider externalnode.TokenProvider, reg *register.DeviceRegistration, ) (*nodev1.ExternalNode, error) { diff --git a/pkg/cmd/allowssh/allowssh_test.go b/pkg/cmd/enablessh/enablessh_test.go similarity index 95% rename from pkg/cmd/allowssh/allowssh_test.go rename to pkg/cmd/enablessh/enablessh_test.go index 42ed275c4..c2c3e753b 100644 --- a/pkg/cmd/allowssh/allowssh_test.go +++ b/pkg/cmd/enablessh/enablessh_test.go @@ -1,4 +1,4 @@ -package allowssh +package enablessh import ( "context" @@ -225,12 +225,12 @@ func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider return register.NewNodeServiceClient(provider, m.serverURL) } -type mockAllowSSHStore struct { +type mockEnableSSHStore struct { token string } -func (m *mockAllowSSHStore) GetCurrentUser() (*entity.User, error) { return &entity.User{}, nil } -func (m *mockAllowSSHStore) GetAccessToken() (string, error) { return m.token, nil } +func (m *mockEnableSSHStore) GetCurrentUser() (*entity.User, error) { return &entity.User{}, nil } +func (m *mockEnableSSHStore) GetAccessToken() (string, error) { return m.token, nil } // mockSelector implements terminal.Selector, returning the first item. type mockSelector struct{ choice string } @@ -280,12 +280,12 @@ func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1 return connect.NewResponse(resp), nil } -func startFakeServer(t *testing.T, svc *fakeNodeService) allowSSHDeps { +func startFakeServer(t *testing.T, svc *fakeNodeService) enableSSHDeps { t.Helper() _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) server := httptest.NewServer(handler) t.Cleanup(server.Close) - return allowSSHDeps{ + return enableSSHDeps{ nodeClients: mockNodeClientFactory{serverURL: server.URL}, prompter: mockSelector{}, } @@ -304,7 +304,7 @@ func Test_fetchRegisteredNode(t *testing.T) { }, } deps := startFakeServer(t, svc) - store := &mockAllowSSHStore{token: "tok"} + store := &mockEnableSSHStore{token: "tok"} reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} node, err := fetchRegisteredNode(context.Background(), deps, store, reg) @@ -380,7 +380,7 @@ func Test_installCertAuthority(t *testing.T) { }) } -func Test_allowSSH_LegacyNodeFallsBackToKeys(t *testing.T) { +func Test_enableSSH_LegacyNodeFallsBackToKeys(t *testing.T) { svc := &fakeNodeService{ getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { return &nodev1.GetNodeResponse{ @@ -415,7 +415,7 @@ func Test_allowSSH_LegacyNodeFallsBackToKeys(t *testing.T) { } term := terminal.New() - if err := allowSSH(context.Background(), term, deps, &mockAllowSSHStore{}, reg); err != nil { + if err := enableSSH(context.Background(), term, deps, &mockEnableSSHStore{}, reg); err != nil { t.Fatalf("allowSSH failed: %v", err) } diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 0619d4f84..2d1a8dbd6 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -89,7 +89,7 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode)") cmd.Flags().StringVarP(&nodeFlag, "node", "n", "", "node name (required in non-interactive mode)") cmd.Flags().StringVarP(&userFlag, "user", "u", "", "Brev user ID or email to grant (required in non-interactive mode)") - cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (e.g. after allow-ssh on the node)") + cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (e.g. after enable-ssh on the node)") cmd.Flags().StringVar(&portIDFlag, "port-id", "", "Brev port ID to grant access on (required in non-interactive mode)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip confirmation prompt (assume yes)") @@ -185,7 +185,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt if len(linuxUserOptions) > 0 { linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions) } else { - // No SSH access entries yet (e.g. first grant after allow-ssh): + // No SSH access entries yet (e.g. first grant after enable-ssh): // the node's Linux users aren't known to Brev, and grant-ssh may // run off-box, so ask for the target Linux user. linuxUser = deps.inputter.Input(terminal.PromptContent{ diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 1fc4bc3eb..66ecabe31 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -87,7 +87,7 @@ var ( registerLong = `Register your device with NVIDIA Brev This command registers this machine with Brev and brings up the Brev tunnel. -Registration does not enable SSH; run 'brev allow-ssh' afterwards to allow SSH +Registration does not enable SSH; run 'brev enable-ssh' afterwards to enable SSH on this device, then 'brev grant-ssh' to grant users SSH access. Two modes are supported: @@ -108,7 +108,7 @@ flow is used.` brev register --name my-node --org my-org # Allow SSH on this device after registering - brev allow-ssh` + brev enable-ssh` ) func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { @@ -141,7 +141,7 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { cmd.Flags().StringVarP(&nameFlag, "name", "n", "", "device name (required when using non-interactive mode)") cmd.Flags().IntVarP(&sshPort, "ssh-port", "p", 0, "SSH port (if ssh access is desired)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip all confirmation prompts (assume yes)") - _ = cmd.Flags().MarkDeprecated("ssh-port", "run 'brev allow-ssh' after registration to allow SSH on this device") + _ = cmd.Flags().MarkDeprecated("ssh-port", "use 'brev enable-ssh' after registration to enable SSH access") return cmd } @@ -363,7 +363,7 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore t.Vprintf("%s Registration complete.\n", t.Green(" ✓")) t.Vprint("") - t.Vprintf(" %s\n", t.Green("To allow SSH on this device, run: brev allow-ssh")) + t.Vprintf(" %s\n", t.Green("To enable SSH access to this device, run: brev enable-ssh")) return nil } From aa4adcbb942dff6f60fd462bc795de2657fa923f Mon Sep 17 00:00:00 2001 From: Drew Malin Date: Tue, 1 Sep 2026 11:38:58 -0700 Subject: [PATCH 4/5] handle existing port allocations but do not assume 22, allow prompts or flags for linux user and ssh port in enable-ssh, enable-ssh and grant-ssh both default to the current linux user but both allow this to be overridden, derive organization from the api key if one is provided --- pkg/cmd/enablessh/enablessh.go | 171 +++++++++++++-- pkg/cmd/enablessh/enablessh_test.go | 319 +++++++++++++++++++++++++++- pkg/cmd/grantssh/grantssh.go | 87 ++++---- pkg/cmd/grantssh/grantssh_test.go | 109 +++++++++- pkg/cmd/register/register.go | 32 +-- pkg/cmd/register/register_test.go | 39 +++- pkg/cmd/register/sshkeys.go | 17 +- 7 files changed, 680 insertions(+), 94 deletions(-) diff --git a/pkg/cmd/enablessh/enablessh.go b/pkg/cmd/enablessh/enablessh.go index 23908d6d1..fea155ccf 100644 --- a/pkg/cmd/enablessh/enablessh.go +++ b/pkg/cmd/enablessh/enablessh.go @@ -34,8 +34,11 @@ type enableSSHDeps struct { nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore prompter terminal.Selector + promptLinuxUser func(*terminal.Terminal, string) (string, error) + promptSSHPort func(*terminal.Terminal) (int32, error) // currentUser resolves the OS user for authorized_keys operations. currentUser func() (*user.User, error) + lookupUser func(string) (*user.User, error) } func defaultEnableSSHDeps() enableSSHDeps { @@ -44,37 +47,59 @@ func defaultEnableSSHDeps() enableSSHDeps { nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), prompter: register.TerminalPrompter{}, + promptLinuxUser: register.PromptLinuxUsername, + promptSSHPort: register.PromptSSHPort, currentUser: user.Current, + lookupUser: user.Lookup, } } func NewCmdEnableSSH(t *terminal.Terminal, store EnableSSHStore) *cobra.Command { + var linuxUserFlag string + var sshPortFlag int32 + cmd := &cobra.Command{ Annotations: map[string]string{"configuration": ""}, Use: "enable-ssh", DisableFlagsInUseLine: true, Short: "Trust the Brev certificate authority on this device for SSH", - Long: "Writes the Brev certificate authority to authorized_keys, allowing this device to be an SSH target for the current Linux user. Users are granted access with 'brev grant-ssh'.", - Example: " brev enable-ssh", + Long: "Writes the Brev certificate authority to authorized_keys, allowing this device to be an SSH target. Interactive mode prompts for the Linux user and SSH port. Non-interactive mode requires both --linux-user and --ssh-port. Users are granted access with 'brev grant-ssh'.", + Example: " # Interactive\n brev enable-ssh\n\n # Non-interactive\n brev enable-ssh --linux-user ubuntu --ssh-port 22", RunE: func(cmd *cobra.Command, args []string) error { - return runEnableSSH(cmd.Context(), t, store, defaultEnableSSHDeps()) + interactive := !cmd.Flags().Changed("linux-user") && !cmd.Flags().Changed("ssh-port") + return runEnableSSH(cmd.Context(), t, store, defaultEnableSSHDeps(), enableSSHOpts{ + interactive: interactive, + linuxUsername: linuxUserFlag, + sshPort: sshPortFlag, + }) }, } + cmd.Flags().StringVar(&linuxUserFlag, "linux-user", "", "Linux username to enable SSH for (required in non-interactive mode)") + cmd.Flags().Int32Var(&sshPortFlag, "ssh-port", 0, "SSH destination port (required in non-interactive mode)") return cmd } -func runEnableSSH(ctx context.Context, t *terminal.Terminal, s EnableSSHStore, deps enableSSHDeps) error { +type enableSSHOpts struct { + interactive bool + linuxUsername string + sshPort int32 +} + +func runEnableSSH(ctx context.Context, t *terminal.Terminal, s EnableSSHStore, deps enableSSHDeps, opts enableSSHOpts) error { if !deps.platform.IsCompatible() { return fmt.Errorf("brev enable-ssh is only supported on Linux") } + if err := validateEnableSSHOpts(opts); err != nil { + return err + } reg, err := deps.registrationStore.Load() if err != nil { return fmt.Errorf("failed to read registration file: %w", err) } - return enableSSH(ctx, t, deps, s, reg) + return enableSSH(ctx, t, deps, s, reg, opts) } func enableSSH( @@ -83,10 +108,11 @@ func enableSSH( deps enableSSHDeps, s EnableSSHStore, reg *register.DeviceRegistration, + opts enableSSHOpts, ) error { - linuxUser, err := deps.currentUser() + linuxUser, err := resolveLinuxUser(t, deps, opts) if err != nil { - return fmt.Errorf("failed to determine current Linux user: %w", err) + return err } linuxUsername := linuxUser.Username @@ -106,7 +132,7 @@ func enableSSH( return fmt.Errorf("enable SSH failed: %w", err) } if node.GetLabels()[sshcert.LabelKeySSHProvider] != sshcert.SSHProviderCertAuth { - return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername) + return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername, opts) } caPublicKey = node.GetCertificateAuthority() } @@ -116,7 +142,7 @@ func enableSSH( } t.Vprint(t.Green(" Certificate authority written to authorized_keys.")) - if err := ensureSSHPort(ctx, t, deps, s, reg); err != nil { + if err := ensureSSHPort(ctx, t, deps, s, reg, opts); err != nil { return fmt.Errorf("enable SSH failed: %w", err) } @@ -125,37 +151,108 @@ func enableSSH( return nil } -func ensureSSHPort(ctx context.Context, t *terminal.Terminal, deps enableSSHDeps, s EnableSSHStore, reg *register.DeviceRegistration) error { +func ensureSSHPort(ctx context.Context, t *terminal.Terminal, deps enableSSHDeps, s EnableSSHStore, reg *register.DeviceRegistration, opts enableSSHOpts) error { + sshPort := opts.sshPort + if opts.interactive { + var err error + sshPort, err = deps.promptSSHPort(t) + if err != nil { + return fmt.Errorf("reading SSH port: %w", err) + } + } + ports, err := fetchRegisteredNode(ctx, deps, s, reg) if err != nil { t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Note: could not check existing ports: %v", err))) } - if p := findExistingSSHPort(ports); p != nil { + if p := findExistingSSHPort(ports, sshPort); p != nil { t.Vprintf(" SSH port already allocated (%s).\n", register.FormatPortLabel(p)) return nil } - sshPort, err := register.PromptSSHPort(t) - if err != nil { - return fmt.Errorf("reading SSH port: %w", err) - } - if _, err := register.OpenSSHPort(ctx, t, deps.nodeClients, s, reg, sshPort); err != nil { + // A prior invocation or concurrent request may have allocated the port + // even when OpenPort returns an error. Confirm the resulting state before + // deciding whether the operation failed. + refreshed, refreshErr := fetchRegisteredNode(ctx, deps, s, reg) + if refreshErr == nil { + if p := findExistingSSHPort(refreshed, sshPort); p != nil { + t.Vprintf(" SSH port already allocated (%s).\n", register.FormatPortLabel(p)) + return nil + } + } + if isPortAlreadyAllocatedError(err, sshPort) { + t.Vprintf(" SSH port %d already allocated.\n", sshPort) + return nil + } return breverrors.WrapAndTrace(err) } return nil } -func findExistingSSHPort(node *nodev1.ExternalNode) *nodev1.Port { +func findExistingSSHPort(node *nodev1.ExternalNode, destinationPort int32) *nodev1.Port { for _, p := range node.GetPorts() { - if p.GetPortNumber() == 22 { + if p.GetServerPort() == destinationPort || (p.GetServerPort() == 0 && p.GetPortNumber() == destinationPort) { return p } } return nil } +func isPortAlreadyAllocatedError(err error, port int32) bool { + if err == nil { + return false + } + return strings.Contains(strings.ToLower(err.Error()), fmt.Sprintf("port %d is already allocated", port)) +} + +func validateEnableSSHOpts(opts enableSSHOpts) error { + if opts.interactive { + return nil + } + if strings.TrimSpace(opts.linuxUsername) == "" || opts.sshPort == 0 { + return fmt.Errorf("in non-interactive mode --linux-user and --ssh-port are required") + } + if opts.sshPort < 1 || opts.sshPort > 65535 { + return fmt.Errorf("invalid --ssh-port %d: port must be between 1 and 65535", opts.sshPort) + } + return nil +} + +func resolveLinuxUser(t *terminal.Terminal, deps enableSSHDeps, opts enableSSHOpts) (*user.User, error) { + if opts.interactive { + return promptLinuxUser(t, deps) + } + linuxUsername := strings.TrimSpace(opts.linuxUsername) + linuxUser, err := deps.lookupUser(linuxUsername) + if err != nil { + return nil, fmt.Errorf("failed to find Linux user %q: %w", linuxUsername, err) + } + return linuxUser, nil +} + +func promptLinuxUser(t *terminal.Terminal, deps enableSSHDeps) (*user.User, error) { + currentLinuxUser, err := deps.currentUser() + if err != nil { + return nil, fmt.Errorf("failed to determine current Linux user: %w", err) + } + linuxUsername, err := deps.promptLinuxUser(t, currentLinuxUser.Username) + if err != nil { + return nil, fmt.Errorf("reading Linux username: %w", err) + } + linuxUsername = strings.TrimSpace(linuxUsername) + if linuxUsername == currentLinuxUser.Username { + return currentLinuxUser, nil + } + + linuxUser, err := deps.lookupUser(linuxUsername) + if err != nil { + return nil, fmt.Errorf("failed to find Linux user %q: %w", linuxUsername, err) + } + return linuxUser, nil +} + func legacyEnableSSH( ctx context.Context, t *terminal.Terminal, @@ -164,13 +261,14 @@ func legacyEnableSSH( reg *register.DeviceRegistration, node *nodev1.ExternalNode, linuxUsername string, + opts enableSSHOpts, ) error { brevUser, err := s.GetCurrentUser() if err != nil { return fmt.Errorf("enable SSH failed: %w", err) } - brevPortID, err := register.ResolveSSHAccessPort(ctx, t, deps.prompter, deps.nodeClients, s, reg, node) + brevPortID, err := resolveLegacySSHPort(ctx, t, deps, s, reg, node, opts) if err != nil { return fmt.Errorf("enable SSH failed: %w", err) } @@ -184,6 +282,41 @@ func legacyEnableSSH( return nil } +func resolveLegacySSHPort( + ctx context.Context, + t *terminal.Terminal, + deps enableSSHDeps, + s EnableSSHStore, + reg *register.DeviceRegistration, + node *nodev1.ExternalNode, + opts enableSSHOpts, +) (string, error) { + if opts.interactive { + portID, err := register.ResolveSSHAccessPort(ctx, t, deps.prompter, deps.nodeClients, s, reg, node) + if err != nil { + return "", fmt.Errorf("resolve SSH access port: %w", err) + } + return portID, nil + } + if p := findExistingSSHPort(node, opts.sshPort); p != nil { + t.Vprintf(" SSH port already allocated (%s).\n", register.FormatPortLabel(p)) + return p.GetPortId(), nil + } + + portID, err := register.OpenSSHPort(ctx, t, deps.nodeClients, s, reg, opts.sshPort) + if err == nil { + return portID, nil + } + refreshed, refreshErr := fetchRegisteredNode(ctx, deps, s, reg) + if refreshErr == nil { + if p := findExistingSSHPort(refreshed, opts.sshPort); p != nil { + t.Vprintf(" SSH port already allocated (%s).\n", register.FormatPortLabel(p)) + return p.GetPortId(), nil + } + } + return "", fmt.Errorf("open SSH port: %w", err) +} + func installCertAuthority(osUser *user.User, caPublicKey, nodeID, linuxUser string) error { if caPublicKey == "" { return fmt.Errorf("certificate authority public key is required") diff --git a/pkg/cmd/enablessh/enablessh_test.go b/pkg/cmd/enablessh/enablessh_test.go index c2c3e753b..a01f238e5 100644 --- a/pkg/cmd/enablessh/enablessh_test.go +++ b/pkg/cmd/enablessh/enablessh_test.go @@ -252,6 +252,7 @@ func (m mockSelector) Select(_ string, items []string) string { type fakeNodeService struct { nodev1connect.UnimplementedExternalNodeServiceHandler getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) + openPortFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) grantCalls int openCalls int } @@ -263,11 +264,19 @@ func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, _ *connect.Reque func (f *fakeNodeService) OpenPort(_ context.Context, req *connect.Request[nodev1.OpenPortRequest]) (*connect.Response[nodev1.OpenPortResponse], error) { f.openCalls++ + if f.openPortFn != nil { + resp, err := f.openPortFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil + } return connect.NewResponse(&nodev1.OpenPortResponse{ Port: &nodev1.Port{ PortId: fmt.Sprintf("port_%d", req.Msg.GetPortNumber()), Protocol: req.Msg.GetProtocol(), - PortNumber: req.Msg.GetPortNumber(), + PortNumber: 41000, + ServerPort: req.Msg.GetPortNumber(), }, }), nil } @@ -288,6 +297,10 @@ func startFakeServer(t *testing.T, svc *fakeNodeService) enableSSHDeps { return enableSSHDeps{ nodeClients: mockNodeClientFactory{serverURL: server.URL}, prompter: mockSelector{}, + promptLinuxUser: func(_ *terminal.Terminal, defaultUsername string) (string, error) { + return defaultUsername, nil + }, + promptSSHPort: register.PromptSSHPort, } } @@ -316,6 +329,221 @@ func Test_fetchRegisteredNode(t *testing.T) { } } +func Test_findExistingSSHPortMatchesDestinationPort(t *testing.T) { + node := &nodev1.ExternalNode{Ports: []*nodev1.Port{ + {PortId: "port_other", PortNumber: 22, ServerPort: 8080}, + {PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}, + }} + + got := findExistingSSHPort(node, 22) + if got == nil || got.GetPortId() != "port_ssh" { + t.Fatalf("expected destination port 22 mapping, got %+v", got) + } +} + +func Test_isPortAlreadyAllocatedError(t *testing.T) { + err := fmt.Errorf("failed to allocate port: internal: 400 Bad Request: Port 22 is already allocated for this client") + if !isPortAlreadyAllocatedError(err, 22) { + t.Fatal("expected matching already-allocated error to be recognized") + } + if isPortAlreadyAllocatedError(err, 2222) { + t.Fatal("must not recognize an allocation error for a different port") + } + if isPortAlreadyAllocatedError(fmt.Errorf("permission denied"), 22) { + t.Fatal("must not recognize an unrelated error") + } +} + +func Test_ensureSSHPortPromptsThenReusesSelectedPort(t *testing.T) { + register.SetTestSSHPort(22) + defer register.ClearTestSSHPort() + + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{Ports: []*nodev1.Port{ + {PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}, + }}}, nil + }, + } + deps := startFakeServer(t, svc) + reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} + + if err := ensureSSHPort(context.Background(), terminal.New(), deps, &mockEnableSSHStore{}, reg, enableSSHOpts{interactive: true}); err != nil { + t.Fatalf("ensureSSHPort: %v", err) + } + if svc.openCalls != 0 { + t.Fatalf("OpenPort called %d times, want 0", svc.openCalls) + } +} + +func Test_ensureSSHPortDoesNotAssumeExistingDifferentPort(t *testing.T) { + register.SetTestSSHPort(2222) + defer register.ClearTestSSHPort() + + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{Ports: []*nodev1.Port{ + {PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}, + }}}, nil + }, + } + deps := startFakeServer(t, svc) + reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} + + if err := ensureSSHPort(context.Background(), terminal.New(), deps, &mockEnableSSHStore{}, reg, enableSSHOpts{interactive: true}); err != nil { + t.Fatalf("ensureSSHPort: %v", err) + } + if svc.openCalls != 1 { + t.Fatalf("OpenPort called %d times, want 1", svc.openCalls) + } +} + +func Test_ensureSSHPortTreatsCreateErrorAsSuccessWhenPortNowExists(t *testing.T) { + register.SetTestSSHPort(22) + defer register.ClearTestSSHPort() + + getCalls := 0 + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + getCalls++ + node := &nodev1.ExternalNode{} + if getCalls > 1 { + node.Ports = []*nodev1.Port{{PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}} + } + return &nodev1.GetNodeResponse{ExternalNode: node}, nil + }, + openPortFn: func(_ *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) { + return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("Port 22 is already allocated for this client")) + }, + } + deps := startFakeServer(t, svc) + reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} + + if err := ensureSSHPort(context.Background(), terminal.New(), deps, &mockEnableSSHStore{}, reg, enableSSHOpts{interactive: true}); err != nil { + t.Fatalf("ensureSSHPort: %v", err) + } + if svc.openCalls != 1 || getCalls != 2 { + t.Fatalf("got OpenPort calls=%d GetNode calls=%d, want 1 and 2", svc.openCalls, getCalls) + } +} + +func Test_ensureSSHPortNonInteractiveDoesNotPrompt(t *testing.T) { + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{Ports: []*nodev1.Port{ + {PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}, + }}}, nil + }, + } + deps := startFakeServer(t, svc) + deps.promptSSHPort = func(*terminal.Terminal) (int32, error) { + t.Fatal("non-interactive mode must not prompt for the SSH port") + return 0, nil + } + reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} + + err := ensureSSHPort(context.Background(), terminal.New(), deps, &mockEnableSSHStore{}, reg, enableSSHOpts{ + linuxUsername: "ubuntu", + sshPort: 22, + }) + if err != nil { + t.Fatalf("ensureSSHPort: %v", err) + } + if svc.openCalls != 0 { + t.Fatalf("OpenPort called %d times, want 0", svc.openCalls) + } +} + +func Test_validateEnableSSHOpts(t *testing.T) { + tests := []struct { + name string + opts enableSSHOpts + wantErr string + }{ + {name: "interactive", opts: enableSSHOpts{interactive: true}}, + {name: "non-interactive", opts: enableSSHOpts{linuxUsername: "ubuntu", sshPort: 22}}, + {name: "missing user", opts: enableSSHOpts{sshPort: 22}, wantErr: "--linux-user and --ssh-port are required"}, + {name: "missing port", opts: enableSSHOpts{linuxUsername: "ubuntu"}, wantErr: "--linux-user and --ssh-port are required"}, + {name: "invalid port", opts: enableSSHOpts{linuxUsername: "ubuntu", sshPort: 65536}, wantErr: "port must be between 1 and 65535"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateEnableSSHOpts(tt.opts) + if tt.wantErr == "" { + if err != nil { + t.Fatalf("validateEnableSSHOpts: %v", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) + } + }) + } +} + +func Test_promptLinuxUser(t *testing.T) { + current := &user.User{Username: "current", HomeDir: "/current"} + target := &user.User{Username: "ubuntu", HomeDir: "/home/ubuntu"} + promptCalls := 0 + var promptDefault string + deps := enableSSHDeps{ + currentUser: func() (*user.User, error) { return current, nil }, + promptLinuxUser: func(_ *terminal.Terminal, defaultUsername string) (string, error) { + promptCalls++ + promptDefault = defaultUsername + return defaultUsername, nil + }, + lookupUser: func(username string) (*user.User, error) { + if username != "ubuntu" { + t.Fatalf("lookup username = %q, want ubuntu", username) + } + return target, nil + }, + } + + term := terminal.New() + got, err := promptLinuxUser(term, deps) + if err != nil || got != current { + t.Fatalf("default user = %+v, err = %v", got, err) + } + if promptCalls != 1 { + t.Fatalf("interactive prompt calls = %d, want 1", promptCalls) + } + if promptDefault != "current" { + t.Fatalf("prompt default = %q, want current", promptDefault) + } + + deps.promptLinuxUser = func(_ *terminal.Terminal, defaultUsername string) (string, error) { + promptCalls++ + return "ubuntu", nil + } + got, err = promptLinuxUser(term, deps) + if err != nil || got != target { + t.Fatalf("prompted user = %+v, err = %v", got, err) + } + if promptCalls != 2 { + t.Fatalf("interactive prompt calls = %d, want 2", promptCalls) + } + + got, err = resolveLinuxUser(term, deps, enableSSHOpts{linuxUsername: "ubuntu", sshPort: 22}) + if err != nil || got != target { + t.Fatalf("non-interactive user = %+v, err = %v", got, err) + } + if promptCalls != 2 { + t.Fatalf("non-interactive mode should bypass the prompt; calls = %d", promptCalls) + } +} + +func Test_NewCmdEnableSSHExposesNonInteractiveFlags(t *testing.T) { + cmd := NewCmdEnableSSH(terminal.New(), &mockEnableSSHStore{}) + for _, flagName := range []string{"linux-user", "ssh-port"} { + if cmd.Flags().Lookup(flagName) == nil { + t.Fatalf("enable-ssh should expose --%s for non-interactive mode", flagName) + } + } +} + // --- installCertAuthority --- func Test_installCertAuthority(t *testing.T) { @@ -380,6 +608,93 @@ func Test_installCertAuthority(t *testing.T) { }) } +func Test_resolveLegacySSHPortNonInteractiveUsesProvidedPort(t *testing.T) { + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{}}, nil + }, + } + deps := startFakeServer(t, svc) + node := &nodev1.ExternalNode{Ports: []*nodev1.Port{ + {PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}, + }} + reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} + + portID, err := resolveLegacySSHPort( + context.Background(), + terminal.New(), + deps, + &mockEnableSSHStore{}, + reg, + node, + enableSSHOpts{linuxUsername: "ubuntu", sshPort: 22}, + ) + if err != nil { + t.Fatalf("resolveLegacySSHPort: %v", err) + } + if portID != "port_ssh" { + t.Fatalf("port ID = %q, want port_ssh", portID) + } + if svc.openCalls != 0 { + t.Fatalf("OpenPort called %d times, want 0", svc.openCalls) + } +} + +func Test_enableSSHNonInteractiveUsesProvidedInputs(t *testing.T) { + const caKey = "ssh-ed25519 AAAAC3Nz dummyCA" + svc := &fakeNodeService{ + getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{Ports: []*nodev1.Port{ + {PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}, + }}}, nil + }, + } + deps := startFakeServer(t, svc) + targetUser := &user.User{Username: "ubuntu", HomeDir: t.TempDir()} + promptCalls := 0 + deps.promptLinuxUser = func(*terminal.Terminal, string) (string, error) { + promptCalls++ + return "", nil + } + deps.currentUser = func() (*user.User, error) { + t.Fatal("non-interactive mode must not resolve the current Linux user") + return nil, nil + } + deps.lookupUser = func(username string) (*user.User, error) { + if username != "ubuntu" { + t.Fatalf("lookup username = %q, want ubuntu", username) + } + return targetUser, nil + } + deps.promptSSHPort = func(*terminal.Terminal) (int32, error) { + t.Fatal("non-interactive mode must not prompt for the SSH port") + return 0, nil + } + reg := ®ister.DeviceRegistration{ + ExternalNodeID: "unode_abc", + OrgID: "org_1", + CertificateAuthority: caKey, + } + + err := enableSSH(context.Background(), terminal.New(), deps, &mockEnableSSHStore{}, reg, enableSSHOpts{ + linuxUsername: "ubuntu", + sshPort: 22, + }) + if err != nil { + t.Fatalf("enableSSH: %v", err) + } + if promptCalls != 0 { + t.Fatalf("Linux username prompt calls = %d, want 0", promptCalls) + } + if svc.openCalls != 0 { + t.Fatalf("OpenPort called %d times, want 0", svc.openCalls) + } + wantPrincipal := `principals="brev:v1:vm:unode_abc:login:ubuntu"` + if authorizedKeys := readAuthorizedKeys(t, targetUser); !strings.Contains(authorizedKeys, wantPrincipal) { + t.Fatalf("authorized_keys missing %s:\n%s", wantPrincipal, authorizedKeys) + } +} + func Test_enableSSH_LegacyNodeFallsBackToKeys(t *testing.T) { svc := &fakeNodeService{ getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { @@ -415,7 +730,7 @@ func Test_enableSSH_LegacyNodeFallsBackToKeys(t *testing.T) { } term := terminal.New() - if err := enableSSH(context.Background(), term, deps, &mockEnableSSHStore{}, reg); err != nil { + if err := enableSSH(context.Background(), term, deps, &mockEnableSSHStore{}, reg, enableSSHOpts{interactive: true}); err != nil { t.Fatalf("allowSSH failed: %v", err) } diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 2d1a8dbd6..5a346f0c4 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -5,13 +5,13 @@ package grantssh import ( "context" "fmt" - "maps" - "slices" + "os/user" "strings" nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" + "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" @@ -25,6 +25,7 @@ import ( // GrantSSHStore defines the store methods needed by the grant-ssh command. type GrantSSHStore interface { + auth.APIKeyAuthStore GetCurrentUser() (*entity.User, error) GetActiveOrganizationOrDefault() (*entity.Organization, error) GetOrganizationsByName(name string) ([]entity.Organization, error) @@ -38,9 +39,10 @@ type GrantSSHStore interface { // can be replaced in tests. type grantSSHDeps struct { prompter terminal.Selector - inputter terminal.Inputter + promptLinuxUser func(*terminal.Terminal, string) (string, error) nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore + currentUser func() (*user.User, error) } type resolvedMember struct { @@ -50,9 +52,10 @@ type resolvedMember struct { func defaultGrantSSHDeps() grantSSHDeps { return grantSSHDeps{ prompter: register.TerminalPrompter{}, - inputter: register.TerminalPrompter{}, + promptLinuxUser: register.PromptLinuxUsername, nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), + currentUser: user.Current, } } @@ -69,8 +72,8 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { Use: "grant-ssh", DisableFlagsInUseLine: true, Short: "Grant SSH access to a node for another org member", - Long: "Grant SSH access to a node for another member of your organization. Interactive: no flags, prompts for org, node, port, user, and Linux user. Non-interactive: --org, --node, --user, --linux-user, and --port-id required.", - Example: " brev grant-ssh\n brev grant-ssh --org my-org --node my-node --user user@example.com --linux-user ubuntu --port-id port_abc --approve", + Long: "Grant SSH access to a node for another member of your organization. Interactive: no flags, prompts for org (unless API-key auth is active), node, port, user, and Linux user. Non-interactive: --node, --user, and --port-id are required; --org is also required unless API-key auth is active. The Linux user defaults to the current user and can be changed with --linux-user.", + Example: " brev grant-ssh\n brev grant-ssh --org my-org --node my-node --user user@example.com --port-id port_abc --approve\n brev grant-ssh --node my-node --user user@example.com --linux-user ubuntu --port-id port_abc --approve --api-key ", RunE: func(cmd *cobra.Command, args []string) error { interactive := orgFlag == "" && nodeFlag == "" && userFlag == "" && linuxUser == "" && portIDFlag == "" opts := grantSSHOpts{ @@ -86,10 +89,10 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { }, } - cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode)") + cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode unless using API-key auth)") cmd.Flags().StringVarP(&nodeFlag, "node", "n", "", "node name (required in non-interactive mode)") cmd.Flags().StringVarP(&userFlag, "user", "u", "", "Brev user ID or email to grant (required in non-interactive mode)") - cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (e.g. after enable-ssh on the node)") + cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (defaults to the current user)") cmd.Flags().StringVar(&portIDFlag, "port-id", "", "Brev port ID to grant access on (required in non-interactive mode)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip confirmation prompt (assume yes)") @@ -109,12 +112,13 @@ type grantSSHOpts struct { // runGrantSSH runs the grant-ssh flow; the only difference by mode is whether we prompt or use opts. func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opts grantSSHOpts, deps grantSSHDeps) error { //nolint:gocognit,gocyclo,funlen // ok + apiKeyAuth := auth.IsAPIKeyAuthStore(s) if !opts.interactive { - if opts.orgName == "" || opts.nodeName == "" || opts.userIDOrEmail == "" { - return fmt.Errorf("in non-interactive mode --org, --node, and --user are required") + if opts.nodeName == "" || opts.userIDOrEmail == "" { + return fmt.Errorf("in non-interactive mode --node and --user are required") } - if opts.linuxUser == "" { - return fmt.Errorf("--linux-user is required in non-interactive mode") + if opts.orgName == "" && !apiKeyAuth { + return fmt.Errorf("in non-interactive mode --org is required unless using API-key auth") } if opts.portID == "" { return fmt.Errorf("--port-id is required in non-interactive mode") @@ -123,13 +127,16 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt var org *entity.Organization var err error - if opts.interactive { + switch { + case apiKeyAuth: + org, err = register.ResolveOrgForAPIKey(s, opts.orgName) + case opts.interactive: allOrgs, listErr := s.ListOrganizations() if listErr != nil { return breverrors.WrapAndTrace(listErr) } org, err = helpers.SelectOrganizationInteractive(t, allOrgs, deps.prompter) - } else { + default: org, err = helpers.ResolveOrgByName(s, opts.orgName) } if err != nil { @@ -169,7 +176,6 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt } var selectedUser *entity.User - var linuxUser string if opts.interactive { usersToSelect := make([]string, len(orgMembers)) for i, r := range orgMembers { @@ -180,27 +186,24 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt if err != nil { return err } - linuxUserOptions := uniqueLinuxUsersFromNodeSSHAccess(node) - t.Vprint("") - if len(linuxUserOptions) > 0 { - linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions) - } else { - // No SSH access entries yet (e.g. first grant after enable-ssh): - // the node's Linux users aren't known to Brev, and grant-ssh may - // run off-box, so ask for the target Linux user. - linuxUser = deps.inputter.Input(terminal.PromptContent{ - Label: "Linux username on the node", - ErrorMsg: "linux user is required", - AllowEmpty: false, - }) - linuxUser = strings.TrimSpace(linuxUser) - } } else { selectedUser, err = findUserByIDOrEmail(orgMembers, opts.userIDOrEmail) if err != nil { return err } - linuxUser = opts.linuxUser + } + + linuxUser, err := resolveLinuxUsername(opts.linuxUser, deps.currentUser) + if err != nil { + return err + } + if opts.interactive { + t.Vprint("") + linuxUser, err = deps.promptLinuxUser(t, linuxUser) + if err != nil { + return fmt.Errorf("reading Linux username: %w", err) + } + linuxUser = strings.TrimSpace(linuxUser) } t.Vprint("") @@ -272,18 +275,20 @@ func findPortByID(node *nodev1.ExternalNode, portID string) *nodev1.Port { return nil } -// uniqueLinuxUsersFromNodeSSHAccess returns distinct Linux users from the node's existing SSH access (for picker). -func uniqueLinuxUsersFromNodeSSHAccess(node *nodev1.ExternalNode) []string { - if node == nil { - return nil +func resolveLinuxUsername(linuxUsername string, currentUser func() (*user.User, error)) (string, error) { + linuxUsername = strings.TrimSpace(linuxUsername) + if linuxUsername != "" { + return linuxUsername, nil } - linuxUsers := make(map[string]struct{}) - for _, sa := range node.GetSshAccess() { - if u := sa.GetLinuxUser(); u != "" { - linuxUsers[u] = struct{}{} - } + linuxUser, err := currentUser() + if err != nil { + return "", fmt.Errorf("failed to determine current Linux user: %w", err) + } + linuxUsername = strings.TrimSpace(linuxUser.Username) + if linuxUsername == "" { + return "", fmt.Errorf("failed to determine current Linux user: username is empty") } - return slices.Collect(maps.Keys(linuxUsers)) + return linuxUsername, nil } func findUserByIDOrEmail(members []resolvedMember, idOrEmail string) (*entity.User, error) { diff --git a/pkg/cmd/grantssh/grantssh_test.go b/pkg/cmd/grantssh/grantssh_test.go index 4910aefba..3e4b34624 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "net/http/httptest" + "os/user" + "strings" "testing" nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" @@ -71,9 +73,12 @@ type mockGrantSSHStore struct { token string members []*nodev1.OrganizationMember users map[string]*entity.User + tokens *entity.AuthTokens err error } +func (m *mockGrantSSHStore) GetAuthTokens() (*entity.AuthTokens, error) { return m.tokens, nil } + func (m *mockGrantSSHStore) GetCurrentUser() (*entity.User, error) { if m.err != nil { return nil, m.err @@ -179,17 +184,15 @@ func testGrantSSHDeps(t *testing.T, svc *fakeNodeService, regStore register.Regi } return "" }}, - inputter: mockInputter{value: "testuser"}, + promptLinuxUser: func(_ *terminal.Terminal, defaultUsername string) (string, error) { + return defaultUsername, nil + }, nodeClients: mockNodeClientFactory{serverURL: server.URL}, registrationStore: regStore, + currentUser: func() (*user.User, error) { return &user.User{Username: "ubuntu"}, nil }, }, server } -// mockInputter implements terminal.Inputter, returning a fixed value. -type mockInputter struct{ value string } - -func (m mockInputter) Input(_ terminal.PromptContent) string { return m.value } - func Test_runGrantSSH_NotRegistered(t *testing.T) { regStore := &mockRegistrationStore{} // no registration @@ -275,6 +278,11 @@ func Test_runGrantSSH_HappyPath(t *testing.T) { deps, server := testGrantSSHDeps(t, svc, regStore) defer server.Close() + var linuxUserDefault string + deps.promptLinuxUser = func(_ *terminal.Terminal, defaultUsername string) (string, error) { + linuxUserDefault = defaultUsername + return "dmalin", nil + } term := terminal.New() opts := grantSSHOpts{interactive: true, skipConfirm: true, linuxUser: "ubuntu"} @@ -292,8 +300,11 @@ func Test_runGrantSSH_HappyPath(t *testing.T) { if gotReq.GetUserId() != "user_2" { t.Errorf("expected user ID user_2, got %s", gotReq.GetUserId()) } - if gotReq.GetLinuxUser() != "ubuntu" { - t.Errorf("expected linux user ubuntu, got %s", gotReq.GetLinuxUser()) + if gotReq.GetLinuxUser() != "dmalin" { + t.Errorf("expected selected Linux user dmalin, got %s", gotReq.GetLinuxUser()) + } + if linuxUserDefault != "ubuntu" { + t.Errorf("Linux username prompt default = %q, want ubuntu", linuxUserDefault) } if gotReq.GetPortId() != "port_ssh" { t.Errorf("expected port ID port_ssh, got %s", gotReq.GetPortId()) @@ -352,7 +363,6 @@ func Test_runGrantSSH_NonInteractiveWithPortID(t *testing.T) { orgName: "TestOrg", nodeName: "My Spark", userIDOrEmail: "alice@example.com", - linuxUser: "ubuntu", portID: "port_ssh", skipConfirm: true, } @@ -363,6 +373,87 @@ func Test_runGrantSSH_NonInteractiveWithPortID(t *testing.T) { if gotReq == nil || gotReq.GetPortId() != "port_ssh" { t.Fatalf("expected port_ssh in request, got %+v", gotReq) } + if gotReq.GetLinuxUser() != "ubuntu" { + t.Fatalf("expected current Linux user ubuntu, got %q", gotReq.GetLinuxUser()) + } +} + +func Test_runGrantSSH_APIKeyDerivesOrganization(t *testing.T) { + targetUser := &entity.User{ID: "user_2", Name: "Alice", Email: "alice@example.com"} + store := &mockGrantSSHStore{ + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + tokens: &entity.AuthTokens{APIKey: "bak-test", APIKeyOrgID: "org_123"}, + members: []*nodev1.OrganizationMember{ + {UserId: "user_2"}, + }, + users: map[string]*entity.User{"user_2": targetUser}, + } + + var gotListReq *nodev1.ListNodesRequest + var gotGrantReq *nodev1.GrantNodeSSHAccessRequest + svc := &fakeNodeService{ + listNodesFn: func(req *nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) { + gotListReq = req + return &nodev1.ListNodesResponse{Items: []*nodev1.ExternalNode{{ + ExternalNodeId: "unode_abc", + Name: "My Spark", + Ports: []*nodev1.Port{{PortId: "port_ssh", PortNumber: 11640, ServerPort: 22}}, + }}}, nil + }, + grantSSHFn: func(req *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) { + gotGrantReq = req + return &nodev1.GrantNodeSSHAccessResponse{}, nil + }, + } + deps, server := testGrantSSHDeps(t, svc, &mockRegistrationStore{}) + defer server.Close() + + err := runGrantSSH(context.Background(), terminal.New(), store, grantSSHOpts{ + interactive: false, + nodeName: "My Spark", + userIDOrEmail: "alice@example.com", + portID: "port_ssh", + skipConfirm: true, + }, deps) + if err != nil { + t.Fatalf("runGrantSSH: %v", err) + } + if gotListReq == nil || gotListReq.GetOrganizationId() != "org_123" { + t.Fatalf("ListNodes organization = %+v, want org_123", gotListReq) + } + if gotGrantReq == nil || gotGrantReq.GetLinuxUser() != "ubuntu" { + t.Fatalf("GrantNodeSSHAccess request = %+v, want Linux user ubuntu", gotGrantReq) + } +} + +func Test_runGrantSSH_APIKeyRejectsDifferentOrganization(t *testing.T) { + store := &mockGrantSSHStore{ + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + tokens: &entity.AuthTokens{APIKey: "bak-test", APIKeyOrgID: "org_123"}, + } + err := runGrantSSH(context.Background(), terminal.New(), store, grantSSHOpts{ + interactive: false, + orgName: "OtherOrg", + nodeName: "My Spark", + userIDOrEmail: "alice@example.com", + portID: "port_ssh", + }, grantSSHDeps{}) + if err == nil || !strings.Contains(err.Error(), "api key does not belong to organization") { + t.Fatalf("expected API-key org mismatch error, got %v", err) + } +} + +func Test_resolveLinuxUsername(t *testing.T) { + currentUser := func() (*user.User, error) { return &user.User{Username: "dmalin"}, nil } + + got, err := resolveLinuxUsername("", currentUser) + if err != nil || got != "dmalin" { + t.Fatalf("default username = %q, err = %v", got, err) + } + got, err = resolveLinuxUsername(" ubuntu ", currentUser) + if err != nil || got != "ubuntu" { + t.Fatalf("explicit username = %q, err = %v", got, err) + } } func Test_runGrantSSH_RPCFailure(t *testing.T) { diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 66ecabe31..662098591 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -28,6 +28,7 @@ import ( // RegisterStore defines the store methods needed by the register command. type RegisterStore interface { + auth.APIKeyAuthStore GetCurrentUser() (*entity.User, error) GetActiveOrganizationOrDefault() (*entity.Organization, error) GetOrganizationsByName(name string) ([]entity.Organization, error) @@ -93,20 +94,22 @@ on this device, then 'brev grant-ssh' to grant users SSH access. Two modes are supported: • Interactive (default): run 'brev register' with no flags and follow prompts for device name and org. • Non-interactive: use --name and --org. No prompts; --name is required, and - --org is required unless --api-key is supplied. Use for scripts/CI. + --org is required unless API-key auth is active. Use for scripts/CI. -Headless auth (credential chain): pass --api-key (a Brev API key) or set -the BREV_API_KEY environment variable to authenticate without the login -link; the key authenticates this register command only — run 'brev login ---api-key' afterward to stay logged in. If neither is set, the login-link -flow is used.` +API-key auth: pass --api-key, set BREV_API_KEY, or first run 'brev login +--api-key'. A key passed directly authenticates this register command only; +run 'brev login --api-key' to save it. If no API-key auth is active, the +login-link flow is used.` registerExample = ` # Interactive (prompts for device name, org, confirmations) brev register - # Non-interactive (--name and --org required) + # Non-interactive with user auth brev register --name my-node --org my-org + # Non-interactive with API-key auth (the org is derived from the key) + brev register --name my-node --api-key + # Allow SSH on this device after registering brev enable-ssh` ) @@ -137,7 +140,7 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { }, } - cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required when using non-interactive mode)") + cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode unless using API-key auth)") cmd.Flags().StringVarP(&nameFlag, "name", "n", "", "device name (required when using non-interactive mode)") cmd.Flags().IntVarP(&sshPort, "ssh-port", "p", 0, "SSH port (if ssh access is desired)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip all confirmation prompts (assume yes)") @@ -169,22 +172,23 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt return breverrors.NewValidationError(fmt.Sprintf("api key must be a Brev API key (expected %s prefix); see 'brev login --api-key'", auth.BrevAPIKeyPrefix)) } } + apiKeyAuth := apiKey != "" || auth.IsAPIKeyAuthStore(s) if !opts.interactive { if opts.name == "" { return fmt.Errorf("in non-interactive mode --name is required") } - if opts.orgName == "" && apiKey == "" { - return fmt.Errorf("in non-interactive mode --org is required unless --api-key is supplied") + if opts.orgName == "" && !apiKeyAuth { + return fmt.Errorf("in non-interactive mode --org is required unless using API-key auth") } } - if err := isAuthenticated(s, apiKey); err != nil { + if err := isAuthenticated(s, apiKeyAuth); err != nil { return breverrors.WrapAndTrace(err) } var intendedOrg *entity.Organization switch { - case apiKey != "": + case apiKeyAuth: o, err := ResolveOrgForAPIKey(s, opts.orgName) if err != nil { return err @@ -379,8 +383,8 @@ func resolveOrgInteractive(t *terminal.Terminal, s RegisterStore, deps registerD return org, nil } -func isAuthenticated(s RegisterStore, apiKey string) error { - if apiKey != "" { +func isAuthenticated(s RegisterStore, apiKeyAuth bool) error { + if apiKeyAuth { return nil } if _, err := s.GetCurrentUser(); err != nil { diff --git a/pkg/cmd/register/register_test.go b/pkg/cmd/register/register_test.go index 0124bfb2b..68a7a95b8 100644 --- a/pkg/cmd/register/register_test.go +++ b/pkg/cmd/register/register_test.go @@ -21,13 +21,16 @@ import ( // mockRegisterStore satisfies RegisterStore for orchestration tests. type mockRegisterStore struct { - user *entity.User - org *entity.Organization - orgs []entity.Organization - token string - err error + user *entity.User + org *entity.Organization + orgs []entity.Organization + token string + tokens *entity.AuthTokens + err error } +func (m *mockRegisterStore) GetAuthTokens() (*entity.AuthTokens, error) { return m.tokens, nil } + func (m *mockRegisterStore) GetCurrentUser() (*entity.User, error) { if m.err != nil { return nil, m.err @@ -847,6 +850,26 @@ func Test_runRegister_APIKey(t *testing.T) { } } +func Test_runRegister_SavedAPIKeyUsesKeyOrganization(t *testing.T) { + ensureNoAPIKeyEnv(t) + regStore := &mockRegistrationStore{} + var gotOrgs []string + svc := &fakeNodeService{addNodeFn: okAddNodeFn(&gotOrgs)} + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + store := testRegisterStore() + store.tokens = &entity.AuthTokens{APIKey: testAPIKey, APIKeyOrgID: "org_123"} + + err := runRegister(context.Background(), terminal.New(), store, + registerOpts{interactive: false, name: "my-spark"}, deps) + if err != nil { + t.Fatalf("runRegister: %v", err) + } + if len(gotOrgs) != 1 || gotOrgs[0] != "org_123" { + t.Fatalf("expected API-key org org_123, got %v", gotOrgs) + } +} + func Test_runRegister_OrgMismatch(t *testing.T) { tests := []struct { name string @@ -938,21 +961,21 @@ func Test_isAuthenticated(t *testing.T) { store := &mockRegisterStore{ err: fmt.Errorf("GetCurrentUser must not be called when a key is present"), } - if err := isAuthenticated(store, testAPIKey); err != nil { + if err := isAuthenticated(store, true); err != nil { t.Fatalf("expected nil error with api key, got %v", err) } }) t.Run("no api key verifies via GetCurrentUser", func(t *testing.T) { store := &mockRegisterStore{user: &entity.User{ID: "user_1"}} - if err := isAuthenticated(store, ""); err != nil { + if err := isAuthenticated(store, false); err != nil { t.Fatalf("expected nil error with valid user, got %v", err) } }) t.Run("no api key and GetCurrentUser fails", func(t *testing.T) { store := &mockRegisterStore{err: fmt.Errorf("not logged in")} - err := isAuthenticated(store, "") + err := isAuthenticated(store, false) if err == nil { t.Fatal("expected error when GetCurrentUser fails") } diff --git a/pkg/cmd/register/sshkeys.go b/pkg/cmd/register/sshkeys.go index 1188766dc..443d1b7b3 100644 --- a/pkg/cmd/register/sshkeys.go +++ b/pkg/cmd/register/sshkeys.go @@ -273,7 +273,6 @@ func openPortForSSHAccess( } // OpenSSHPort calls the OpenPort RPC to allocate a port on the node for SSH access. -// The call is idempotent — if the port is already open, the server returns the existing allocation. func OpenSSHPort( ctx context.Context, t *terminal.Terminal, @@ -390,6 +389,22 @@ func SetTestSSHPort(port int32) { testSSHPort = &port } // ClearTestSSHPort clears the test port override. func ClearTestSSHPort() { testSSHPort = nil } +// PromptLinuxUsername prompts for a Linux username and returns defaultUsername +// when the user presses Enter without typing a value. +func PromptLinuxUsername(t *terminal.Terminal, defaultUsername string) (string, error) { + t.Vprintf(" %s ", t.Green(fmt.Sprintf("Linux username (default %s):", defaultUsername))) + reader := bufio.NewReader(os.Stdin) + line, err := reader.ReadString('\n') + if err != nil { + return "", fmt.Errorf("reading input: %w", err) + } + linuxUsername := strings.TrimSpace(line) + if linuxUsername == "" { + return defaultUsername, nil + } + return linuxUsername, nil +} + // PromptSSHPort prompts the user for the target SSH port, defaulting to 22 if // they press Enter or leave it empty. Re-prompts on invalid input until a valid // port is provided. Only returns an error for unrecoverable I/O failures. From ef0333fe4cc02ba05117f96bc2677a9cccc86f61 Mon Sep 17 00:00:00 2001 From: Drew Malin Date: Tue, 1 Sep 2026 11:57:43 -0700 Subject: [PATCH 5/5] simplify enableSSHDeps struct for testing --- pkg/cmd/enablessh/enablessh.go | 10 +-- pkg/cmd/enablessh/enablessh_test.go | 74 ++++++++++--------- pkg/cmd/grantssh/grantssh.go | 6 +- pkg/cmd/grantssh/grantssh_test.go | 35 ++++++--- pkg/cmd/register/providers.go | 11 +++ pkg/cmd/register/sshkeys.go | 46 +++++------- pkg/cmd/register/sshkeys_port_resolve_test.go | 15 ++-- pkg/cmd/register/sshkeys_test.go | 36 ++++++++- pkg/terminal/types.go | 5 ++ 9 files changed, 141 insertions(+), 97 deletions(-) diff --git a/pkg/cmd/enablessh/enablessh.go b/pkg/cmd/enablessh/enablessh.go index fea155ccf..a5f1b1387 100644 --- a/pkg/cmd/enablessh/enablessh.go +++ b/pkg/cmd/enablessh/enablessh.go @@ -33,9 +33,7 @@ type enableSSHDeps struct { platform externalnode.PlatformChecker nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore - prompter terminal.Selector - promptLinuxUser func(*terminal.Terminal, string) (string, error) - promptSSHPort func(*terminal.Terminal) (int32, error) + prompter register.SSHAccessPrompter // currentUser resolves the OS user for authorized_keys operations. currentUser func() (*user.User, error) lookupUser func(string) (*user.User, error) @@ -47,8 +45,6 @@ func defaultEnableSSHDeps() enableSSHDeps { nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), prompter: register.TerminalPrompter{}, - promptLinuxUser: register.PromptLinuxUsername, - promptSSHPort: register.PromptSSHPort, currentUser: user.Current, lookupUser: user.Lookup, } @@ -155,7 +151,7 @@ func ensureSSHPort(ctx context.Context, t *terminal.Terminal, deps enableSSHDeps sshPort := opts.sshPort if opts.interactive { var err error - sshPort, err = deps.promptSSHPort(t) + sshPort, err = register.PromptSSHPort(t, deps.prompter) if err != nil { return fmt.Errorf("reading SSH port: %w", err) } @@ -237,7 +233,7 @@ func promptLinuxUser(t *terminal.Terminal, deps enableSSHDeps) (*user.User, erro if err != nil { return nil, fmt.Errorf("failed to determine current Linux user: %w", err) } - linuxUsername, err := deps.promptLinuxUser(t, currentLinuxUser.Username) + linuxUsername, err := register.PromptLinuxUsername(t, deps.prompter, currentLinuxUser.Username) if err != nil { return nil, fmt.Errorf("reading Linux username: %w", err) } diff --git a/pkg/cmd/enablessh/enablessh_test.go b/pkg/cmd/enablessh/enablessh_test.go index a01f238e5..859e91a23 100644 --- a/pkg/cmd/enablessh/enablessh_test.go +++ b/pkg/cmd/enablessh/enablessh_test.go @@ -232,8 +232,10 @@ type mockEnableSSHStore struct { func (m *mockEnableSSHStore) GetCurrentUser() (*entity.User, error) { return &entity.User{}, nil } func (m *mockEnableSSHStore) GetAccessToken() (string, error) { return m.token, nil } -// mockSelector implements terminal.Selector, returning the first item. -type mockSelector struct{ choice string } +type mockSelector struct { + choice string + inputLineFunc func(*terminal.Terminal, string) (string, error) +} func (m mockSelector) Select(_ string, items []string) string { if m.choice != "" { @@ -249,6 +251,13 @@ func (m mockSelector) Select(_ string, items []string) string { return "" } +func (m mockSelector) InputLine(t *terminal.Terminal, label string) (string, error) { + if m.inputLineFunc != nil { + return m.inputLineFunc(t, label) + } + return "", nil +} + type fakeNodeService struct { nodev1connect.UnimplementedExternalNodeServiceHandler getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) @@ -297,10 +306,6 @@ func startFakeServer(t *testing.T, svc *fakeNodeService) enableSSHDeps { return enableSSHDeps{ nodeClients: mockNodeClientFactory{serverURL: server.URL}, prompter: mockSelector{}, - promptLinuxUser: func(_ *terminal.Terminal, defaultUsername string) (string, error) { - return defaultUsername, nil - }, - promptSSHPort: register.PromptSSHPort, } } @@ -355,9 +360,6 @@ func Test_isPortAlreadyAllocatedError(t *testing.T) { } func Test_ensureSSHPortPromptsThenReusesSelectedPort(t *testing.T) { - register.SetTestSSHPort(22) - defer register.ClearTestSSHPort() - svc := &fakeNodeService{ getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{Ports: []*nodev1.Port{ @@ -377,9 +379,6 @@ func Test_ensureSSHPortPromptsThenReusesSelectedPort(t *testing.T) { } func Test_ensureSSHPortDoesNotAssumeExistingDifferentPort(t *testing.T) { - register.SetTestSSHPort(2222) - defer register.ClearTestSSHPort() - svc := &fakeNodeService{ getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { return &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{Ports: []*nodev1.Port{ @@ -388,6 +387,9 @@ func Test_ensureSSHPortDoesNotAssumeExistingDifferentPort(t *testing.T) { }, } deps := startFakeServer(t, svc) + deps.prompter = mockSelector{inputLineFunc: func(*terminal.Terminal, string) (string, error) { + return "2222", nil + }} reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} if err := ensureSSHPort(context.Background(), terminal.New(), deps, &mockEnableSSHStore{}, reg, enableSSHOpts{interactive: true}); err != nil { @@ -399,9 +401,6 @@ func Test_ensureSSHPortDoesNotAssumeExistingDifferentPort(t *testing.T) { } func Test_ensureSSHPortTreatsCreateErrorAsSuccessWhenPortNowExists(t *testing.T) { - register.SetTestSSHPort(22) - defer register.ClearTestSSHPort() - getCalls := 0 svc := &fakeNodeService{ getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { @@ -436,9 +435,11 @@ func Test_ensureSSHPortNonInteractiveDoesNotPrompt(t *testing.T) { }, } deps := startFakeServer(t, svc) - deps.promptSSHPort = func(*terminal.Terminal) (int32, error) { - t.Fatal("non-interactive mode must not prompt for the SSH port") - return 0, nil + deps.prompter = mockSelector{ + inputLineFunc: func(*terminal.Terminal, string) (string, error) { + t.Fatal("non-interactive mode must not prompt for input") + return "", nil + }, } reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} @@ -486,13 +487,15 @@ func Test_promptLinuxUser(t *testing.T) { current := &user.User{Username: "current", HomeDir: "/current"} target := &user.User{Username: "ubuntu", HomeDir: "/home/ubuntu"} promptCalls := 0 - var promptDefault string + var promptLabel string deps := enableSSHDeps{ currentUser: func() (*user.User, error) { return current, nil }, - promptLinuxUser: func(_ *terminal.Terminal, defaultUsername string) (string, error) { - promptCalls++ - promptDefault = defaultUsername - return defaultUsername, nil + prompter: mockSelector{ + inputLineFunc: func(_ *terminal.Terminal, label string) (string, error) { + promptCalls++ + promptLabel = label + return "", nil + }, }, lookupUser: func(username string) (*user.User, error) { if username != "ubuntu" { @@ -510,13 +513,15 @@ func Test_promptLinuxUser(t *testing.T) { if promptCalls != 1 { t.Fatalf("interactive prompt calls = %d, want 1", promptCalls) } - if promptDefault != "current" { - t.Fatalf("prompt default = %q, want current", promptDefault) + if promptLabel != "Linux username (default current):" { + t.Fatalf("prompt label = %q", promptLabel) } - deps.promptLinuxUser = func(_ *terminal.Terminal, defaultUsername string) (string, error) { - promptCalls++ - return "ubuntu", nil + deps.prompter = mockSelector{ + inputLineFunc: func(_ *terminal.Terminal, _ string) (string, error) { + promptCalls++ + return "ubuntu", nil + }, } got, err = promptLinuxUser(term, deps) if err != nil || got != target { @@ -652,9 +657,12 @@ func Test_enableSSHNonInteractiveUsesProvidedInputs(t *testing.T) { deps := startFakeServer(t, svc) targetUser := &user.User{Username: "ubuntu", HomeDir: t.TempDir()} promptCalls := 0 - deps.promptLinuxUser = func(*terminal.Terminal, string) (string, error) { - promptCalls++ - return "", nil + deps.prompter = mockSelector{ + inputLineFunc: func(*terminal.Terminal, string) (string, error) { + promptCalls++ + t.Fatal("non-interactive mode must not prompt for input") + return "", nil + }, } deps.currentUser = func() (*user.User, error) { t.Fatal("non-interactive mode must not resolve the current Linux user") @@ -666,10 +674,6 @@ func Test_enableSSHNonInteractiveUsesProvidedInputs(t *testing.T) { } return targetUser, nil } - deps.promptSSHPort = func(*terminal.Terminal) (int32, error) { - t.Fatal("non-interactive mode must not prompt for the SSH port") - return 0, nil - } reg := ®ister.DeviceRegistration{ ExternalNodeID: "unode_abc", OrgID: "org_1", diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 5a346f0c4..5c7ac2433 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -38,8 +38,7 @@ type GrantSSHStore interface { // grantSSHDeps bundles the side-effecting dependencies of runGrantSSH so they // can be replaced in tests. type grantSSHDeps struct { - prompter terminal.Selector - promptLinuxUser func(*terminal.Terminal, string) (string, error) + prompter register.SSHAccessPrompter nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore currentUser func() (*user.User, error) @@ -52,7 +51,6 @@ type resolvedMember struct { func defaultGrantSSHDeps() grantSSHDeps { return grantSSHDeps{ prompter: register.TerminalPrompter{}, - promptLinuxUser: register.PromptLinuxUsername, nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), currentUser: user.Current, @@ -199,7 +197,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt } if opts.interactive { t.Vprint("") - linuxUser, err = deps.promptLinuxUser(t, linuxUser) + linuxUser, err = register.PromptLinuxUsername(t, deps.prompter, linuxUser) if err != nil { return fmt.Errorf("reading Linux username: %w", err) } diff --git a/pkg/cmd/grantssh/grantssh_test.go b/pkg/cmd/grantssh/grantssh_test.go index 3e4b34624..1b2f6aa2e 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -21,11 +21,25 @@ import ( // mock types for grantSSHDeps interfaces type mockSelector struct { - fn func(label string, items []string) string + fn func(label string, items []string) string + inputLineFunc func(*terminal.Terminal, string) (string, error) } func (m mockSelector) Select(label string, items []string) string { - return m.fn(label, items) + if m.fn != nil { + return m.fn(label, items) + } + if len(items) > 0 { + return items[0] + } + return "" +} + +func (m mockSelector) InputLine(t *terminal.Terminal, label string) (string, error) { + if m.inputLineFunc != nil { + return m.inputLineFunc(t, label) + } + return "", nil } type mockNodeClientFactory struct { @@ -184,9 +198,6 @@ func testGrantSSHDeps(t *testing.T, svc *fakeNodeService, regStore register.Regi } return "" }}, - promptLinuxUser: func(_ *terminal.Terminal, defaultUsername string) (string, error) { - return defaultUsername, nil - }, nodeClients: mockNodeClientFactory{serverURL: server.URL}, registrationStore: regStore, currentUser: func() (*user.User, error) { return &user.User{Username: "ubuntu"}, nil }, @@ -278,10 +289,12 @@ func Test_runGrantSSH_HappyPath(t *testing.T) { deps, server := testGrantSSHDeps(t, svc, regStore) defer server.Close() - var linuxUserDefault string - deps.promptLinuxUser = func(_ *terminal.Terminal, defaultUsername string) (string, error) { - linuxUserDefault = defaultUsername - return "dmalin", nil + var linuxUserPromptLabel string + deps.prompter = mockSelector{ + inputLineFunc: func(_ *terminal.Terminal, label string) (string, error) { + linuxUserPromptLabel = label + return "dmalin", nil + }, } term := terminal.New() @@ -303,8 +316,8 @@ func Test_runGrantSSH_HappyPath(t *testing.T) { if gotReq.GetLinuxUser() != "dmalin" { t.Errorf("expected selected Linux user dmalin, got %s", gotReq.GetLinuxUser()) } - if linuxUserDefault != "ubuntu" { - t.Errorf("Linux username prompt default = %q, want ubuntu", linuxUserDefault) + if linuxUserPromptLabel != "Linux username (default ubuntu):" { + t.Errorf("Linux username prompt label = %q", linuxUserPromptLabel) } if gotReq.GetPortId() != "port_ssh" { t.Errorf("expected port ID port_ssh, got %s", gotReq.GetPortId()) diff --git a/pkg/cmd/register/providers.go b/pkg/cmd/register/providers.go index 0c6adf843..6970b80e1 100644 --- a/pkg/cmd/register/providers.go +++ b/pkg/cmd/register/providers.go @@ -1,7 +1,9 @@ package register import ( + "bufio" "fmt" + "os" "os/exec" "runtime" "strings" @@ -35,6 +37,15 @@ func (TerminalPrompter) Select(label string, items []string) string { }) } +func (TerminalPrompter) InputLine(t *terminal.Terminal, label string) (string, error) { + t.Vprintf(" %s ", t.Green(label)) + line, err := bufio.NewReader(os.Stdin).ReadString('\n') + if err != nil { + return "", fmt.Errorf("reading input: %w", err) + } + return line, nil +} + // Input prompts for free-form text input. func (TerminalPrompter) Input(pc terminal.PromptContent) string { return terminal.PromptGetInput(pc) diff --git a/pkg/cmd/register/sshkeys.go b/pkg/cmd/register/sshkeys.go index 443d1b7b3..dd35c02da 100644 --- a/pkg/cmd/register/sshkeys.go +++ b/pkg/cmd/register/sshkeys.go @@ -1,7 +1,6 @@ package register import ( - "bufio" "context" "errors" "fmt" @@ -83,11 +82,17 @@ const ( PortChoiceOpenNew = "Open a new port" ) +// SSHAccessPrompter supports SSH selection menus and plain-line input. +type SSHAccessPrompter interface { + terminal.Selector + terminal.LineInputter +} + // ResolveSSHAccessPort prompts for an existing or new port and returns its Brev port ID. func ResolveSSHAccessPort( ctx context.Context, t *terminal.Terminal, - prompter terminal.Selector, + prompter SSHAccessPrompter, nodeClients externalnode.NodeClientFactory, tokenProvider externalnode.TokenProvider, reg *DeviceRegistration, @@ -95,7 +100,7 @@ func ResolveSSHAccessPort( ) (string, error) { ports := node.GetPorts() if len(ports) == 0 { - return openPortForSSHAccess(ctx, t, nodeClients, tokenProvider, reg) + return openPortForSSHAccess(ctx, t, prompter, nodeClients, tokenProvider, reg) } t.Vprint("") @@ -109,7 +114,7 @@ func ResolveSSHAccessPort( t.Vprintf(" Using port %s.\n", FormatPortLabel(selected)) return selected.GetPortId(), nil case PortChoiceOpenNew: - return openPortForSSHAccess(ctx, t, nodeClients, tokenProvider, reg) + return openPortForSSHAccess(ctx, t, prompter, nodeClients, tokenProvider, reg) default: return "", fmt.Errorf("invalid port choice %q", choice) } @@ -260,12 +265,13 @@ func RemoveAuthorizedKeyLine(u *user.User, line string) error { func openPortForSSHAccess( ctx context.Context, t *terminal.Terminal, + inputter terminal.LineInputter, nodeClients externalnode.NodeClientFactory, tokenProvider externalnode.TokenProvider, reg *DeviceRegistration, ) (string, error) { t.Vprint("") - port, err := PromptSSHPort(t) + port, err := PromptSSHPort(t, inputter) if err != nil { return "", fmt.Errorf("invalid port: %w", err) } @@ -378,25 +384,12 @@ func SetupAndRegisterNodeSSHAccess( const defaultSSHPort = 22 -// testSSHPort is set by tests to avoid blocking on stdin. When non-nil, -// PromptSSHPort returns this value without prompting. -var testSSHPort *int32 - -// SetTestSSHPort sets the port returned by PromptSSHPort without prompting. -// Only for use in tests; call ClearTestSSHPort when done. -func SetTestSSHPort(port int32) { testSSHPort = &port } - -// ClearTestSSHPort clears the test port override. -func ClearTestSSHPort() { testSSHPort = nil } - // PromptLinuxUsername prompts for a Linux username and returns defaultUsername // when the user presses Enter without typing a value. -func PromptLinuxUsername(t *terminal.Terminal, defaultUsername string) (string, error) { - t.Vprintf(" %s ", t.Green(fmt.Sprintf("Linux username (default %s):", defaultUsername))) - reader := bufio.NewReader(os.Stdin) - line, err := reader.ReadString('\n') +func PromptLinuxUsername(t *terminal.Terminal, inputter terminal.LineInputter, defaultUsername string) (string, error) { + line, err := inputter.InputLine(t, fmt.Sprintf("Linux username (default %s):", defaultUsername)) if err != nil { - return "", fmt.Errorf("reading input: %w", err) + return "", fmt.Errorf("prompting for Linux username: %w", err) } linuxUsername := strings.TrimSpace(line) if linuxUsername == "" { @@ -408,16 +401,11 @@ func PromptLinuxUsername(t *terminal.Terminal, defaultUsername string) (string, // PromptSSHPort prompts the user for the target SSH port, defaulting to 22 if // they press Enter or leave it empty. Re-prompts on invalid input until a valid // port is provided. Only returns an error for unrecoverable I/O failures. -func PromptSSHPort(t *terminal.Terminal) (int32, error) { - if testSSHPort != nil { - return *testSSHPort, nil - } - reader := bufio.NewReader(os.Stdin) +func PromptSSHPort(t *terminal.Terminal, inputter terminal.LineInputter) (int32, error) { for { - t.Vprintf(" %s ", t.Green("SSH port (default 22):")) - line, err := reader.ReadString('\n') + line, err := inputter.InputLine(t, "SSH port (default 22):") if err != nil { - return 0, fmt.Errorf("reading input: %w", err) + return 0, fmt.Errorf("prompting for SSH port: %w", err) } portStr := strings.TrimSpace(line) if portStr == "" { diff --git a/pkg/cmd/register/sshkeys_port_resolve_test.go b/pkg/cmd/register/sshkeys_port_resolve_test.go index 243342015..9751876ba 100644 --- a/pkg/cmd/register/sshkeys_port_resolve_test.go +++ b/pkg/cmd/register/sshkeys_port_resolve_test.go @@ -14,6 +14,11 @@ import ( type mockPortSelector struct { choices []string idx int + input string +} + +func (m *mockPortSelector) InputLine(*terminal.Terminal, string) (string, error) { + return m.input, nil } func (m *mockPortSelector) Select(_ string, items []string) string { @@ -59,14 +64,11 @@ func startPortOpenTestServer(t *testing.T) (mockNodeClientFactory, *openPortCapt } func TestResolveSSHAccessPort_noPortsOpensNew(t *testing.T) { - SetTestSSHPort(2222) - defer ClearTestSSHPort() - clients, cap := startPortOpenTestServer(t) portID, err := ResolveSSHAccessPort( context.Background(), terminal.New(), - &mockPortSelector{}, + &mockPortSelector{input: "2222"}, clients, portOpenTestStore{token: "tok"}, &DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"}, @@ -111,11 +113,8 @@ func TestResolveSSHAccessPort_useExisting(t *testing.T) { } func TestResolveSSHAccessPort_openNewWhenPortsExist(t *testing.T) { - SetTestSSHPort(2222) - defer ClearTestSSHPort() - clients, cap := startPortOpenTestServer(t) - sel := &mockPortSelector{choices: []string{PortChoiceOpenNew}} + sel := &mockPortSelector{choices: []string{PortChoiceOpenNew}, input: "2222"} portID, err := ResolveSSHAccessPort( context.Background(), terminal.New(), diff --git a/pkg/cmd/register/sshkeys_test.go b/pkg/cmd/register/sshkeys_test.go index acecb74ab..8983dbbd7 100644 --- a/pkg/cmd/register/sshkeys_test.go +++ b/pkg/cmd/register/sshkeys_test.go @@ -257,14 +257,44 @@ func TestRemoveAuthorizedKey_ByPublicKeyMaterial(t *testing.T) { // --- PromptSSHPort --- +type mockLineInputter struct { + inputs []string + labels []string +} + +func (m *mockLineInputter) InputLine(_ *terminal.Terminal, label string) (string, error) { + m.labels = append(m.labels, label) + input := m.inputs[0] + m.inputs = m.inputs[1:] + return input, nil +} + func TestPromptSSHPort(t *testing.T) { - SetTestSSHPort(2222) - defer ClearTestSSHPort() - port, err := PromptSSHPort(terminal.New()) + inputter := &mockLineInputter{inputs: []string{"not-a-port", "2222"}} + port, err := PromptSSHPort(terminal.New(), inputter) if err != nil { t.Fatalf("PromptSSHPort: %v", err) } if port != 2222 { t.Errorf("expected 2222, got %d", port) } + if len(inputter.labels) != 2 || inputter.labels[0] != "SSH port (default 22):" { + t.Errorf("unexpected prompt labels: %v", inputter.labels) + } +} + +func TestPromptLinuxUsername(t *testing.T) { + inputter := &mockLineInputter{inputs: []string{"", "ubuntu"}} + + username, err := PromptLinuxUsername(terminal.New(), inputter, "dmalin") + if err != nil || username != "dmalin" { + t.Fatalf("default username = %q, err = %v", username, err) + } + username, err = PromptLinuxUsername(terminal.New(), inputter, "dmalin") + if err != nil || username != "ubuntu" { + t.Fatalf("entered username = %q, err = %v", username, err) + } + if len(inputter.labels) != 2 || inputter.labels[0] != "Linux username (default dmalin):" { + t.Errorf("unexpected prompt labels: %v", inputter.labels) + } } diff --git a/pkg/terminal/types.go b/pkg/terminal/types.go index d686723aa..e5ac2d981 100644 --- a/pkg/terminal/types.go +++ b/pkg/terminal/types.go @@ -14,3 +14,8 @@ type Selector interface { type Inputter interface { Input(pc PromptContent) string } + +// LineInputter reads one line without an interactive terminal UI. +type LineInputter interface { + InputLine(t *Terminal, label string) (string, error) +}