tests

Author: patelspratik

pkg/cmd/deregister/deregister.go

@@ -4,12 +4,14 @@ package deregister
 import (
 	"context"
 	"fmt"
+	"os/user"
 	"runtime"
 
 	nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect"
 	nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
 	"connectrpc.com/connect"
 
+	"github.com/brevdev/brev-cli/pkg/cmd/enablessh"
 	"github.com/brevdev/brev-cli/pkg/cmd/register"
 	"github.com/brevdev/brev-cli/pkg/config"
 	"github.com/brevdev/brev-cli/pkg/entity"
@@ -34,6 +36,7 @@ type deregisterDeps struct {
 	uninstallNetbird  func() error
 	newNodeClient     func(provider register.TokenProvider, baseURL string) nodev1connect.ExternalNodeServiceClient
 	registrationStore register.RegistrationStore
+	removeBrevKeys    func(*user.User) error
 }
 
 func prodDeregisterDeps(brevHome string) deregisterDeps {
@@ -48,6 +51,7 @@ func prodDeregisterDeps(brevHome string) deregisterDeps {
 		uninstallNetbird:  register.UninstallNetbird,
 		newNodeClient:     register.NewNodeServiceClient,
 		registrationStore: register.NewFileRegistrationStore(brevHome),
+		removeBrevKeys:    enablessh.RemoveBrevAuthorizedKeys,
 	}
 }
 
@@ -125,6 +129,19 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore,
 	t.Vprint(t.Green("  Node removed from Brev."))
 	t.Vprint("")
 
+	// Remove Brev SSH keys from authorized_keys.
+	u, err := user.Current()
+	if err != nil {
+		t.Vprintf("  Warning: could not determine current user for SSH key cleanup: %v\n", err)
+	} else {
+		if err := deps.removeBrevKeys(u); err != nil {
+			t.Vprintf("  Warning: failed to remove Brev SSH keys: %v\n", err)
+		} else {
+			t.Vprint(t.Green("  Brev SSH keys removed from authorized_keys."))
+		}
+	}
+	t.Vprint("")
+
 	removeNetbird := deps.promptSelect(
 		"Would you also like to uninstall NetBird?",
 		[]string{"Yes, uninstall NetBird", "No, keep NetBird installed"},

pkg/cmd/deregister/deregister_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http/httptest"
+	"os/user"
 	"testing"
 
 	nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect"
@@ -94,6 +95,7 @@ func testDeregisterDeps(t *testing.T, svc *fakeNodeService, regStore register.Re
 			return register.NewNodeServiceClient(provider, server.URL)
 		},
 		registrationStore: regStore,
+		removeBrevKeys:    func(_ *user.User) error { return nil },
 	}, server
 }
 
@@ -295,3 +297,86 @@ func Test_runDeregister_SkipsNetbirdUninstall(t *testing.T) {
 		t.Error("NetBird uninstall should not be called when user declines")
 	}
 }
+
+func Test_runDeregister_CallsRemoveBrevKeys(t *testing.T) {
+	regStore := &mockRegistrationStore{
+		reg: &register.DeviceRegistration{
+			ExternalNodeID: "unode_abc",
+			DisplayName:    "My Spark",
+			OrgID:          "org_123",
+		},
+	}
+
+	store := &mockDeregisterStore{
+		user:  &entity.User{ID: "user_1"},
+		home:  "/home/testuser/.brev",
+		token: "tok",
+	}
+
+	svc := &fakeNodeService{
+		removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) {
+			return &nodev1.RemoveNodeResponse{}, nil
+		},
+	}
+
+	removeBrevKeysCalled := false
+	deps, server := testDeregisterDeps(t, svc, regStore)
+	defer server.Close()
+	deps.removeBrevKeys = func(_ *user.User) error {
+		removeBrevKeysCalled = true
+		return nil
+	}
+
+	term := terminal.New()
+	err := runDeregister(context.Background(), term, store, deps)
+	if err != nil {
+		t.Fatalf("runDeregister failed: %v", err)
+	}
+
+	if !removeBrevKeysCalled {
+		t.Error("expected removeBrevKeys to be called during deregistration")
+	}
+}
+
+func Test_runDeregister_RemoveBrevKeysFailureIsNonFatal(t *testing.T) {
+	regStore := &mockRegistrationStore{
+		reg: &register.DeviceRegistration{
+			ExternalNodeID: "unode_abc",
+			DisplayName:    "My Spark",
+			OrgID:          "org_123",
+		},
+	}
+
+	store := &mockDeregisterStore{
+		user:  &entity.User{ID: "user_1"},
+		home:  "/home/testuser/.brev",
+		token: "tok",
+	}
+
+	svc := &fakeNodeService{
+		removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) {
+			return &nodev1.RemoveNodeResponse{}, nil
+		},
+	}
+
+	deps, server := testDeregisterDeps(t, svc, regStore)
+	defer server.Close()
+	deps.removeBrevKeys = func(_ *user.User) error {
+		return fmt.Errorf("permission denied")
+	}
+
+	term := terminal.New()
+	err := runDeregister(context.Background(), term, store, deps)
+	if err != nil {
+		t.Fatalf("expected deregister to succeed despite removeBrevKeys failure, got: %v", err)
+	}
+
+	// Registration should still be cleaned up.
+	exists, err := regStore.Exists()
+	if err != nil {
+		t.Fatalf("Exists error: %v", err)
+	}
+	if exists {
+		t.Error("expected registration to be deleted even when SSH key cleanup fails")
+	}
+}

pkg/cmd/enablessh/enablessh.go

@@ -142,8 +142,13 @@ func EnableSSH(
 	return nil
 }
 
+// BrevKeyComment is the marker appended to every SSH key that Brev installs.
+// It allows RemoveBrevAuthorizedKeys to identify and remove exactly those keys.
+const BrevKeyComment = "# brev-cli"
+
 // InstallAuthorizedKey appends the given public key to the user's
-// ~/.ssh/authorized_keys if it isn't already present.
+// ~/.ssh/authorized_keys if it isn't already present. The key is tagged with
+// a brev-cli comment so it can be removed later by RemoveBrevAuthorizedKeys.
 func InstallAuthorizedKey(u *user.User, pubKey string) error {
 	pubKey = strings.TrimSpace(pubKey)
 	if pubKey == "" {
@@ -163,15 +168,17 @@ func InstallAuthorizedKey(u *user.User, pubKey string) error {
 	}
 
 	if strings.Contains(string(existing), pubKey) {
-		return nil // already present
+		return nil // already present (tagged or not)
 	}
 
+	taggedKey := pubKey + " " + BrevKeyComment
+
 	// Ensure existing content ends with a newline before appending.
 	content := string(existing)
 	if len(content) > 0 && !strings.HasSuffix(content, "\n") {
 		content += "\n"
 	}
-	content += pubKey + "\n"
+	content += taggedKey + "\n"
 
 	if err := os.WriteFile(authKeysPath, []byte(content), 0o600); err != nil {
 		return fmt.Errorf("writing authorized_keys: %w", err)
@@ -180,6 +187,34 @@ func InstallAuthorizedKey(u *user.User, pubKey string) error {
 	return nil
 }
 
+// RemoveBrevAuthorizedKeys removes all SSH keys tagged with the brev-cli
+// comment from the user's ~/.ssh/authorized_keys.
+func RemoveBrevAuthorizedKeys(u *user.User) error {
+	authKeysPath := filepath.Join(u.HomeDir, ".ssh", "authorized_keys")
+
+	existing, err := os.ReadFile(authKeysPath) // #nosec G304
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return fmt.Errorf("reading authorized_keys: %w", err)
+	}
+
+	var kept []string
+	for _, line := range strings.Split(string(existing), "\n") {
+		if strings.Contains(line, BrevKeyComment) {
+			continue
+		}
+		kept = append(kept, line)
+	}
+
+	result := strings.Join(kept, "\n")
+	if err := os.WriteFile(authKeysPath, []byte(result), 0o600); err != nil {
+		return fmt.Errorf("writing authorized_keys: %w", err)
+	}
+	return 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) {