Add --current-device capabilities to tsh (#30636) (#30702)

* Add the `tsh device asset-tag` hidden command

* Implement registration in the fake device service

* Add `--current-device` to `tsh device enroll`

lib/devicetrust/authn/authn_test.go

@@ -28,7 +28,9 @@ import (
 )
 
 func TestRunCeremony(t *testing.T) {
-	env := testenv.MustNew()
+	env := testenv.MustNew(
+		testenv.WithAutoCreateDevice(true),
+	)
 	defer env.Close()
 
 	devices := env.DevicesClient
@@ -99,7 +101,7 @@ func enrollDevice(ctx context.Context, devices devicepb.DeviceTrustServiceClient
 	if err != nil {
 		return fmt.Errorf("enroll device init: %w", err)
 	}
-	enrollDeviceInit.Token = "fake device token"
+	enrollDeviceInit.Token = testenv.FakeEnrollmentToken
 	if err := stream.Send(&devicepb.EnrollDeviceRequest{
 		Payload: &devicepb.EnrollDeviceRequest_Init{
 			Init: enrollDeviceInit,

lib/devicetrust/enroll/auto_enroll_test.go

@@ -26,7 +26,9 @@ import (
 )
 
 func TestAutoEnrollCeremony_Run(t *testing.T) {
-	env := testenv.MustNew()
+	env := testenv.MustNew(
+		testenv.WithAutoCreateDevice(true),
+	)
 	defer env.Close()
 
 	devices := env.DevicesClient

lib/devicetrust/enroll/enroll.go

@@ -18,6 +18,8 @@ import (
 	"context"
 
 	"github.com/gravitational/trace"
+	"github.com/gravitational/trace/trail"
+	log "github.com/sirupsen/logrus"
 	"golang.org/x/exp/slices"
 
 	devicepb "github.com/gravitational/teleport/api/gen/proto/go/teleport/devicetrust/v1"
@@ -49,10 +51,114 @@ func NewCeremony() *Ceremony {
 	}
 }
 
-// RunCeremony performs the client-side device enrollment ceremony.
-// Equivalent to `NewCeremony().Run()`.
-func RunCeremony(ctx context.Context, devicesClient devicepb.DeviceTrustServiceClient, debug bool, enrollToken string) (*devicepb.Device, error) {
-	return NewCeremony().Run(ctx, devicesClient, debug, enrollToken)
+// RunAdminOutcome is the outcome of [Ceremony.RunAdmin].
+// It is used to communicate the actions performed.
+type RunAdminOutcome int
+
+const (
+	_ RunAdminOutcome = iota // Zero means nothing happened.
+	DeviceEnrolled
+	DeviceRegistered
+	DeviceRegisteredAndEnrolled
+)
+
+// RunAdmin is a more powerful variant of Run: it attempts to register the
+// current device, creates an enrollment token and uses that token to call Run.
+//
+// Must be called by a user capable of performing all actions above, otherwise
+// it fails.
+//
+// Returns the created or enrolled device, an outcome marker and an error. The
+// zero outcome means everything failed.
+//
+// Note that the device may be created and the ceremony can still fail
+// afterwards, causing a return similar to "return dev, DeviceRegistered, err"
+// (where nothing is "nil").
+func (c *Ceremony) RunAdmin(
+	ctx context.Context,
+	devicesClient devicepb.DeviceTrustServiceClient,
+	debug bool,
+) (*devicepb.Device, RunAdminOutcome, error) {
+	// The init message contains the device collected data.
+	init, err := c.EnrollDeviceInit()
+	if err != nil {
+		return nil, 0, trace.Wrap(err)
+	}
+	cdd := init.DeviceData
+	osType := cdd.OsType
+	assetTag := cdd.SerialNumber
+
+	rewordAccessDenied := func(err error, action string) error {
+		if trace.IsAccessDenied(trail.FromGRPC(err)) {
+			log.WithError(err).Debug(
+				"Device Trust: Redacting access denied error with user-friendly message")
+			return trace.AccessDenied(
+				"User does not have permissions to %s. Contact your cluster device administrator.",
+				action,
+			)
+		}
+		return err
+	}
+
+	// Query for current device.
+	findResp, err := devicesClient.FindDevices(ctx, &devicepb.FindDevicesRequest{
+		IdOrTag: assetTag,
+	})
+	if err != nil {
+		return nil, 0, trace.Wrap(rewordAccessDenied(err, "list devices"))
+	}
+	var currentDev *devicepb.Device
+	for _, dev := range findResp.Devices {
+		if dev.OsType == osType {
+			currentDev = dev
+			log.Debugf(
+				"Device Trust: Found device %q/%v, id=%q",
+				currentDev.AssetTag, devicetrust.FriendlyOSType(currentDev.OsType), currentDev.Id,
+			)
+			break
+		}
+	}
+
+	// If missing, create the device.
+	var outcome RunAdminOutcome
+	if currentDev == nil {
+		currentDev, err = devicesClient.CreateDevice(ctx, &devicepb.CreateDeviceRequest{
+			Device: &devicepb.Device{
+				OsType:   osType,
+				AssetTag: assetTag,
+			},
+			CreateEnrollToken: true, // Save an additional RPC.
+		})
+		if err != nil {
+			return nil, outcome, trace.Wrap(rewordAccessDenied(err, "register devices"))
+		}
+		outcome = DeviceRegistered
+	}
+	// From here onwards, always return `currentDev` and `outcome`!
+
+	// If missing, create a new enrollment token.
+	if currentDev.EnrollToken.GetToken() == "" {
+		currentDev.EnrollToken, err = devicesClient.CreateDeviceEnrollToken(ctx, &devicepb.CreateDeviceEnrollTokenRequest{
+			DeviceId: currentDev.Id,
+		})
+		if err != nil {
+			return currentDev, outcome, trace.Wrap(rewordAccessDenied(err, "create device enrollment tokens"))
+		}
+		log.Debugf(
+			"Device Trust: Created enrollment token for device %q/%s",
+			currentDev.AssetTag,
+			devicetrust.FriendlyOSType(currentDev.OsType))
+	}
+	token := currentDev.EnrollToken.GetToken()
+
+	// Then proceed onto enrollment.
+	enrolled, err := c.Run(ctx, devicesClient, debug, token)
+	if err != nil {
+		return enrolled, outcome, trace.Wrap(err)
+	}
+
+	outcome++ // "0" becomes "Enrolled", "Registered" becomes "RegisteredAndEnrolled".
+	return enrolled, outcome, trace.Wrap(err)
 }
 
 // Run performs the client-side device enrollment ceremony.

lib/devicetrust/enroll/enroll_test.go

@@ -27,13 +27,70 @@ import (
 	"github.com/gravitational/teleport/lib/devicetrust/testenv"
 )
 
-func TestCeremony_Run(t *testing.T) {
+func TestCeremony_RunAdmin(t *testing.T) {
 	env := testenv.MustNew()
 	defer env.Close()
 
 	devices := env.DevicesClient
 	ctx := context.Background()
 
+	nonExistingDev, err := testenv.NewFakeMacOSDevice()
+	require.NoError(t, err, "NewFakeMacOSDevice failed")
+
+	registeredDev, err := testenv.NewFakeMacOSDevice()
+	require.NoError(t, err, "NewFakeMacOSDevice failed")
+
+	// Create the device corresponding to registeredDev.
+	_, err = devices.CreateDevice(ctx, &devicepb.CreateDeviceRequest{
+		Device: &devicepb.Device{
+			OsType:   registeredDev.GetDeviceOSType(),
+			AssetTag: registeredDev.SerialNumber,
+		},
+	})
+	require.NoError(t, err, "CreateDevice(registeredDev) failed")
+
+	tests := []struct {
+		name        string
+		dev         testenv.FakeDevice
+		wantOutcome enroll.RunAdminOutcome
+	}{
+		{
+			name:        "non-existing device",
+			dev:         nonExistingDev,
+			wantOutcome: enroll.DeviceRegisteredAndEnrolled,
+		},
+		{
+			name:        "registered device",
+			dev:         registeredDev,
+			wantOutcome: enroll.DeviceEnrolled,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			c := &enroll.Ceremony{
+				GetDeviceOSType:         test.dev.GetDeviceOSType,
+				EnrollDeviceInit:        test.dev.EnrollDeviceInit,
+				SignChallenge:           test.dev.SignChallenge,
+				SolveTPMEnrollChallenge: test.dev.SolveTPMEnrollChallenge,
+			}
+
+			enrolled, outcome, err := c.RunAdmin(ctx, devices, false /* debug */)
+			require.NoError(t, err, "RunAdmin failed")
+			assert.NotNil(t, enrolled, "RunAdmin returned nil device")
+			assert.Equal(t, test.wantOutcome, outcome, "RunAdmin outcome mismatch")
+		})
+	}
+}
+
+func TestCeremony_Run(t *testing.T) {
+	env := testenv.MustNew(
+		testenv.WithAutoCreateDevice(true),
+	)
+	defer env.Close()
+
+	devices := env.DevicesClient
+	ctx := context.Background()
+
 	macOSDev1, err := testenv.NewFakeMacOSDevice()
 	require.NoError(t, err, "NewFakeMacOSDevice failed")
 	windowsDev1 := testenv.NewFakeWindowsDevice()
@@ -90,7 +147,7 @@ func TestCeremony_Run(t *testing.T) {
 				SolveTPMEnrollChallenge: test.dev.SolveTPMEnrollChallenge,
 			}
 
-			got, err := c.Run(ctx, devices, false, "faketoken")
+			got, err := c.Run(ctx, devices, false /* debug */, testenv.FakeEnrollmentToken)
 			test.assertErr(t, err)
 			test.assertGotDevice(t, got)
 		})