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/cmd.go b/pkg/cmd/cmd.go index a845788cd..4eb80b23d 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -16,6 +16,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/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" @@ -334,6 +335,7 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor cmd.AddCommand(deregister.NewCmdDeregister(t, externalNodeCmdStore)) cmd.AddCommand(upgrade.NewCmdUpgrade(t, noLoginCmdStore)) 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/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/disablessh/disablessh.go b/pkg/cmd/disablessh/disablessh.go new file mode 100644 index 000000000..e284adac5 --- /dev/null +++ b/pkg/cmd/disablessh/disablessh.go @@ -0,0 +1,99 @@ +package disablessh + +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 DisableSSHStore interface { + GetCurrentUser() (*entity.User, error) + GetAccessToken() (string, error) +} + +type disableSSHDeps struct { + platform externalnode.PlatformChecker + registrationStore register.RegistrationStore +} + +func defaultDisableSSHDeps() disableSSHDeps { + return disableSSHDeps{ + platform: register.LinuxPlatform{}, + registrationStore: register.NewFileRegistrationStore(), + } +} + +func NewCmdDisableSSH(t *terminal.Terminal, store DisableSSHStore) *cobra.Command { + cmd := &cobra.Command{ + Annotations: map[string]string{"configuration": ""}, + 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 disable-ssh", + RunE: func(cmd *cobra.Command, args []string) error { + return runDisableSSH(cmd.Context(), t, store, defaultDisableSSHDeps()) + }, + } + + return cmd +} + +func runDisableSSH(_ context.Context, t *terminal.Terminal, _ DisableSSHStore, deps disableSSHDeps) error { + if !deps.platform.IsCompatible() { + return fmt.Errorf("brev disable-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("disable 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 disabled. Run 'brev enable-ssh' to re-enable.")) + return nil +} diff --git a/pkg/cmd/enablessh/enablessh.go b/pkg/cmd/enablessh/enablessh.go index 9788b0e6f..a5f1b1387 100644 --- a/pkg/cmd/enablessh/enablessh.go +++ b/pkg/cmd/enablessh/enablessh.go @@ -1,39 +1,42 @@ -// Package enablessh provides the brev enableSSH command for enabling SSH access -// to a registered external node. +// Package enablessh implements brev enable-ssh. package enablessh 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" + breverrors "github.com/brevdev/brev-cli/pkg/errors" "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/sshcert" "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 + prompter register.SSHAccessPrompter + // currentUser resolves the OS user for authorized_keys operations. + currentUser func() (*user.User, error) + lookupUser func(string) (*user.User, error) } func defaultEnableSSHDeps() enableSSHDeps { @@ -42,87 +45,314 @@ func defaultEnableSSHDeps() enableSSHDeps { nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), prompter: register.TerminalPrompter{}, + 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: "Enable SSH access to this registered device", - Long: "Enable SSH access to this registered device for the current Brev user.", - Example: " brev enable-ssh", + 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. 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) } - brevUser, err := s.GetCurrentUser() - if err != nil { - return breverrors.WrapAndTrace(err) - } - - return enableSSH(ctx, t, deps, s, reg, brevUser) + return enableSSH(ctx, t, deps, s, reg, opts) } -// 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, + s EnableSSHStore, reg *register.DeviceRegistration, - brevUser *entity.User, + opts enableSSHOpts, ) error { - linuxUser, err := user.Current() + 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 checkSSHDaemon(t) t.Vprint("") - t.Vprint(t.Green("Enabling SSH access 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(" Brev user: %s\n", brevUser.ID) t.Vprintf(" Linux user: %s\n", linuxUsername) t.Vprint("") - node, err := fetchRegisteredNode(ctx, deps, tokenProvider, reg) + caPublicKey := reg.CertificateAuthority + if caPublicKey == "" { + node, err := fetchRegisteredNode(ctx, deps, s, reg) + if err != nil { + return fmt.Errorf("enable SSH failed: %w", err) + } + if node.GetLabels()[sshcert.LabelKeySSHProvider] != sshcert.SSHProviderCertAuth { + return legacyEnableSSH(ctx, t, deps, s, reg, node, linuxUsername, opts) + } + caPublicKey = node.GetCertificateAuthority() + } + + if err := installCertAuthority(linuxUser, caPublicKey, reg.ExternalNodeID, linuxUsername); err != nil { + 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, opts); err != nil { + return fmt.Errorf("enable SSH failed: %w", err) + } + + t.Vprint("") + 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 enableSSHDeps, s EnableSSHStore, reg *register.DeviceRegistration, opts enableSSHOpts) error { + sshPort := opts.sshPort + if opts.interactive { + var err error + sshPort, err = register.PromptSSHPort(t, deps.prompter) + 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, sshPort); p != nil { + t.Vprintf(" SSH port already allocated (%s).\n", register.FormatPortLabel(p)) + return nil + } + + 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, destinationPort int32) *nodev1.Port { + for _, p := range node.GetPorts() { + 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 := register.PromptLinuxUsername(t, deps.prompter, 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, + deps enableSSHDeps, + s EnableSSHStore, + 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, tokenProvider, reg, node) + brevPortID, err := resolveLegacySSHPort(ctx, t, deps, s, reg, node, opts) 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 { + if err := register.SetupAndRegisterNodeSSHAccess(ctx, t, deps.nodeClients, s, reg, brevUser, linuxUsername, brevPortID); err != nil { return fmt.Errorf("enable 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 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") + } + + 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 enableSSHDeps, @@ -132,7 +362,6 @@ func fetchRegisteredNode( 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) @@ -140,8 +369,6 @@ func fetchRegisteredNode( 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 diff --git a/pkg/cmd/enablessh/enablessh_test.go b/pkg/cmd/enablessh/enablessh_test.go index 7df94144d..859e91a23 100644 --- a/pkg/cmd/enablessh/enablessh_test.go +++ b/pkg/cmd/enablessh/enablessh_test.go @@ -2,6 +2,7 @@ package enablessh 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")) @@ -228,13 +229,65 @@ type mockEnableSSHStore struct { token string } -func (m *mockEnableSSHStore) GetCurrentUser() (interface{}, error) { return nil, nil } -func (m *mockEnableSSHStore) 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 } + +type mockSelector struct { + choice string + inputLineFunc func(*terminal.Terminal, string) (string, error) +} + +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 "" +} + +func (m mockSelector) InputLine(t *terminal.Terminal, label string) (string, error) { + if m.inputLineFunc != nil { + return m.inputLineFunc(t, label) + } + return "", nil +} -// 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) + openPortFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, 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++ + 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: 41000, + ServerPort: req.Msg.GetPortNumber(), + }, + }), nil } func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) { @@ -245,14 +298,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) enableSSHDeps { t.Helper() _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) server := httptest.NewServer(handler) t.Cleanup(server.Close) return enableSSHDeps{ nodeClients: mockNodeClientFactory{serverURL: server.URL}, - }, server + prompter: mockSelector{}, + } } func Test_fetchRegisteredNode(t *testing.T) { @@ -267,7 +321,7 @@ func Test_fetchRegisteredNode(t *testing.T) { }}, nil }, } - deps, _ := startFakeServer(t, svc) + deps := startFakeServer(t, svc) store := &mockEnableSSHStore{token: "tok"} reg := ®ister.DeviceRegistration{ExternalNodeID: "unode_abc", OrgID: "org_1"} @@ -279,3 +333,422 @@ func Test_fetchRegisteredNode(t *testing.T) { t.Fatalf("unexpected node: %+v", node) } } + +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) { + 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) { + 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.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 { + 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) { + 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.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"} + + 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 promptLabel string + deps := enableSSHDeps{ + currentUser: func() (*user.User, error) { return current, 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" { + 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 promptLabel != "Linux username (default current):" { + t.Fatalf("prompt label = %q", promptLabel) + } + + deps.prompter = mockSelector{ + inputLineFunc: func(_ *terminal.Terminal, _ 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) { + 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_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.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") + 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 + } + 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) { + 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 := enableSSH(context.Background(), term, deps, &mockEnableSSHStore{}, reg, enableSSHOpts{interactive: true}); 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/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 8d49a274a..5c7ac2433 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) @@ -37,9 +38,10 @@ type GrantSSHStore interface { // grantSSHDeps bundles the side-effecting dependencies of runGrantSSH so they // can be replaced in tests. type grantSSHDeps struct { - prompter terminal.Selector + prompter register.SSHAccessPrompter nodeClients externalnode.NodeClientFactory registrationStore register.RegistrationStore + currentUser func() (*user.User, error) } type resolvedMember struct { @@ -51,6 +53,7 @@ func defaultGrantSSHDeps() grantSSHDeps { prompter: register.TerminalPrompter{}, nodeClients: register.DefaultNodeClientFactory{}, registrationStore: register.NewFileRegistrationStore(), + currentUser: user.Current, } } @@ -67,10 +70,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.", - 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 == "" + interactive := orgFlag == "" && nodeFlag == "" && userFlag == "" && linuxUser == "" && portIDFlag == "" opts := grantSSHOpts{ interactive: interactive, orgName: orgFlag, @@ -84,10 +87,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 (required in non-interactive mode)") + 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)") @@ -107,17 +110,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 - currentUser, err := s.GetCurrentUser() - if err != nil { - return breverrors.WrapAndTrace(err) - } - + 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") @@ -125,13 +124,17 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt } var org *entity.Organization - if opts.interactive { + var err error + 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 { @@ -160,7 +163,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 } @@ -171,7 +174,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 { @@ -182,18 +184,24 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt if err != nil { 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 enable-ssh on the node)") - } - t.Vprint("") - linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions) } 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 = register.PromptLinuxUsername(t, deps.prompter, linuxUser) + if err != nil { + return fmt.Errorf("reading Linux username: %w", err) + } + linuxUser = strings.TrimSpace(linuxUser) } t.Vprint("") @@ -265,18 +273,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) { @@ -289,24 +299,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..1b2f6aa2e 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" @@ -19,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 { @@ -71,9 +87,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 @@ -181,6 +200,7 @@ func testGrantSSHDeps(t *testing.T, svc *fakeNodeService, regStore register.Regi }}, nodeClients: mockNodeClientFactory{serverURL: server.URL}, registrationStore: regStore, + currentUser: func() (*user.User, error) { return &user.User{Username: "ubuntu"}, nil }, }, server } @@ -269,6 +289,13 @@ func Test_runGrantSSH_HappyPath(t *testing.T) { deps, server := testGrantSSHDeps(t, svc, regStore) defer server.Close() + var linuxUserPromptLabel string + deps.prompter = mockSelector{ + inputLineFunc: func(_ *terminal.Terminal, label string) (string, error) { + linuxUserPromptLabel = label + return "dmalin", nil + }, + } term := terminal.New() opts := grantSSHOpts{interactive: true, skipConfirm: true, linuxUser: "ubuntu"} @@ -286,8 +313,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 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()) @@ -346,7 +376,6 @@ func Test_runGrantSSH_NonInteractiveWithPortID(t *testing.T) { orgName: "TestOrg", nodeName: "My Spark", userIDOrEmail: "alice@example.com", - linuxUser: "ubuntu", portID: "port_ssh", skipConfirm: true, } @@ -357,6 +386,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) { @@ -397,7 +507,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 +516,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..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,20 @@ 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) +} + // Netbird handles NetBird installation and uninstallation. type Netbird struct{} diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index e1a9366e4..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) @@ -87,23 +88,30 @@ 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 enable-ssh' afterwards to enable 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. • 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) - brev register --name my-node --org my-org` + # 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` ) func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { @@ -132,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)") @@ -164,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 @@ -319,6 +328,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 @@ -332,14 +342,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("") @@ -372,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..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) } @@ -273,7 +279,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, @@ -379,30 +384,28 @@ 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, 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("prompting for Linux username: %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. -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/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) + } +} diff --git a/pkg/terminal/types.go b/pkg/terminal/types.go index e7138531e..e5ac2d981 100644 --- a/pkg/terminal/types.go +++ b/pkg/terminal/types.go @@ -9,3 +9,13 @@ 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 +} + +// LineInputter reads one line without an interactive terminal UI. +type LineInputter interface { + InputLine(t *Terminal, label string) (string, error) +}