src/go.chromium.org/tast-tests/cros/remote/bundles/cros/policy/zero_touch_enrollment.go

// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package policy

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/golang/protobuf/ptypes/empty"
	"go.chromium.org/tast-tests/cros/common/fixture"
	"go.chromium.org/tast-tests/cros/common/policy"
	"go.chromium.org/tast-tests/cros/common/tape"
	"go.chromium.org/tast-tests/cros/remote/gaiaenrollment"
	"go.chromium.org/tast-tests/cros/remote/log"
	"go.chromium.org/tast-tests/cros/services/cros/graphics"
	"go.chromium.org/tast-tests/cros/services/cros/hwsec"
	ps "go.chromium.org/tast-tests/cros/services/cros/policy"
	"go.chromium.org/tast/core/ctxutil"
	"go.chromium.org/tast/core/errors"
	"go.chromium.org/tast/core/exec"
	"go.chromium.org/tast/core/rpc"
	"go.chromium.org/tast/core/ssh"
	"go.chromium.org/tast/core/testing"
)

const zeroTouchEnrollmentTimeout = 30 * time.Minute

func init() {
	testing.AddTest(&testing.Test{
		Func:         ZeroTouchEnrollment,
		LacrosStatus: testing.LacrosVariantUnneeded,
		Desc:         "ZTE Enroll a device without checking policies",
		Contacts: []string{
			"chromeos-commercial-remote-management@google.com",
			"vsavu@google.com",
		},
		BugComponent: "b:1111632",
		Attr:         []string{"group:dmserver-zteenrollment-daily"},
		SoftwareDeps: []string{"reboot", "chrome"},
		ServiceDeps:  []string{"tast.cros.policy.PolicyService", "tast.cros.tape.Service", "tast.cros.hwsec.OwnershipService", "tast.cros.graphics.ScreenshotService"},
		Fixture:      fixture.CleanOwnership,
		Timeout:      zeroTouchEnrollmentTimeout,
		SearchFlags: []*testing.StringPair{{
			Key: "feature_id",
			// Zero Touch Enrollment.
			Value: "screenplay-6f0905f0-9ecd-4974-b4a1-7e4b828b5dc2",
		}, {
			Key: "feature_id",
			// Enroll an unmanaged device that was pre-provisioned for ZTE to
			// ensure it enrolls without user interaction.
			// COM_FOUND_CUJ25_TASK4_WF1
			Value: "screenplay-cd82fc31-3640-4ccb-ba06-33ddffa54733",
		}},
		Params: []testing.Param{
			{
				Name: "autopush",
				Val: gaiaenrollment.TestParams{
					DMServer:             policy.DMServerAlphaURL,
					PoolID:               tape.ZTETestAutomation,
					SerialNumber:         "policy.ZeroTouchEnrollment.serial_number",
					HardwareModel:        "policy.ZeroTouchEnrollment.hardware_model",
					DeviceProvisionToken: "policy.ZeroTouchEnrollment.device_provision_token",
					CustomerID:           "policy.ZeroTouchEnrollment.customer_id",
					BatchKey:             "policy.ZeroTouchEnrollment.batch_key",
				}, // TODO b/346725308 Refactor to use utility and known dependency list.
				ExtraSearchFlags: []*testing.StringPair{{
					Key: "external_dependency", Value: "DMServerAlpha",
				}},
			},
		},
		Vars: []string{
			"ui.signinProfileTestExtensionManifestKey",
			tape.ServiceAccountVar,
			"policy.ZeroTouchEnrollment.serial_number",
			"policy.ZeroTouchEnrollment.hardware_model",
			"policy.ZeroTouchEnrollment.device_provision_token",
			"policy.ZeroTouchEnrollment.customer_id",
			"policy.ZeroTouchEnrollment.batch_key",
		},
	})
}

func ZeroTouchEnrollment(ctx context.Context, s *testing.State) {
	param := s.Param().(gaiaenrollment.TestParams)
	dmServerURL := param.DMServer
	poolID := param.PoolID
	serialNumber := s.RequiredVar(param.SerialNumber)
	hardwareModel := s.RequiredVar(param.HardwareModel)
	deviceProvisionToken := s.RequiredVar(param.DeviceProvisionToken)
	customerID := s.RequiredVar(param.CustomerID)
	batchKey := s.RequiredVar(param.BatchKey)
	enrolled := false

	// Shorten deadline to leave time separately for logging and cleanup.
	cleanupCtx := ctx
	ctx, cancel := ctxutil.Shorten(cleanupCtx, 20*time.Second)
	defer cancel()
	defer log.Collect(cleanupCtx, s.DUT())

	name, err := preProvisionDevice(ctx, serialNumber, hardwareModel, deviceProvisionToken, customerID, batchKey)
	if err != nil {
		s.Fatal("Failed to pre-provision device: ", err)
	}
	defer func(ctx context.Context) {
		if !enrolled {
			if err := deletePreProvisioningRecord(ctx, name, batchKey); err != nil {
				s.Log("Failed to delete pre-provisioning record: ", err)
			}
		}
	}(cleanupCtx)

	if err := setVpdValuesForInitialEnrollment(ctx, s.DUT().Conn()); err != nil {
		s.Fatal("Failed to get VPD ready for ZTE: ", err)
	}

	cl, err := rpc.Dial(ctx, s.DUT(), s.RPCHint())
	if err != nil {
		s.Fatal("Failed to connect to the RPC service on the DUT: ", err)
	}
	defer cl.Close(cleanupCtx)

	defer func() {
		if !s.HasError() {
			return
		}
		captureScreenshot(cleanupCtx, "test-failure", s)
	}()

	tapeClient, err := tape.NewClient(ctx, []byte(s.RequiredVar(tape.ServiceAccountVar)))
	if err != nil {
		s.Fatal("Failed to create tape client: ", err)
	}
	timeout := int32(zeroTouchEnrollmentTimeout.Seconds())
	// Create an account manager and lease a test account for the duration of the test.
	accManager, acc, err := tape.NewOwnedTestAccountManagerFromClient(ctx, tapeClient, false /*lock*/, tape.WithTimeout(timeout), tape.WithPoolID(poolID))
	if err != nil {
		s.Fatal("Failed to create an account manager and lease an account: ", err)
	}
	defer accManager.CleanUp(cleanupCtx)

	// Deprovision the DUT at the end of the test. As devices might get
	// provisioned even when the enrollment fails we need to defer the
	// deprovisioning before enrolling.
	defer func(ctx context.Context) {
		cl, err := rpc.Dial(ctx, s.DUT(), s.RPCHint())
		if err != nil {
			s.Fatal("Failed to connect to the device before deprovisioning: ", err)
		}
		if err := tapeClient.DeprovisionHelper(ctx, cl, acc.OrgUnitPath); err != nil {
			s.Fatal("Failed to deprovision device: ", err)
		}
	}(cleanupCtx)

	// It may take a while for our preprovisioning command to succeed, wait for a bit and then retry ZTE a few times.
	// GoBigSleepLint: Waiting a bit speeds up the test because provisioning takes time.
	testing.Sleep(ctx, time.Minute)
	attempt := 0
	if err := testing.Poll(ctx, func(ctx context.Context) error {
		// Give ZTE attempt 5 minutes to succeed, then retry.
		oobeCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
		defer cancel()

		// Reconnect to the device because cleaning the TPM restarts Chrome.
		attempt++
		cl, err = rpc.Dial(oobeCtx, s.DUT(), s.RPCHint())
		if err != nil {
			return errors.Wrap(err, "failed to connect to the RPC service on the DUT")
		}
		defer cl.Close(ctx)

		policyClient := ps.NewPolicyServiceClient(cl.Conn)

		if _, err := policyClient.ZeroTouchEnrollUsingChrome(oobeCtx, &ps.ZeroTouchEnrollUsingChromeRequest{
			DmserverURL: dmServerURL,
			ManifestKey: s.RequiredVar("ui.signinProfileTestExtensionManifestKey"),
		}); err != nil {
			s.Logf("Failed to ZTE-enroll on attempt %d with error: %s", attempt, err)
			captureScreenshot(ctx, fmt.Sprintf("attempt-%d-failure", attempt), s)

			// Clear TPM to reset any state left by the failed ZTE attempt.
			ownershipClient := hwsec.NewOwnershipServiceClient(cl.Conn)

			if _, err := ownershipClient.EnsureTPMAndSystemStateAreReset(ctx, &empty.Empty{}); err != nil {
				return errors.Wrap(err, "failed to reset the TPM locally")
			}

			return err
		}

		enrolled = true
		return nil
	}, &testing.PollOptions{Interval: 2 * time.Minute}); err != nil {
		s.Fatal("Failed to ZTE enroll using chrome: ", err)
	}
}

func captureScreenshot(ctx context.Context, filename string, s *testing.State) {
	cl, err := rpc.Dial(ctx, s.DUT(), s.RPCHint())
	if err != nil {
		s.Error("Failed to connect to the device to take screenshot: ", err)
		return
	}
	defer cl.Close(ctx)
	screenshotService := graphics.NewScreenshotServiceClient(cl.Conn)
	screenshotService.CaptureScreenshot(ctx, &graphics.CaptureScreenshotRequest{FilePrefix: filename})
}

func setVpdValuesForInitialEnrollment(ctx context.Context, dutConn *ssh.Conn) error {
	// Delete check_enrollment from RW_VPD if it exists.
	if err := dutConn.CommandContext(ctx, "vpd", "-i", "RW_VPD", "-g", "check_enrollment").Run(); err == nil {
		if err := dutConn.CommandContext(ctx, "vpd", "-i", "RW_VPD", "-d", "check_enrollment").Run(exec.DumpLogOnError); err != nil {
			return errors.Wrap(err, "failed to delete check_enrollment")
		}
	}

	// Delete ActivationDate from RW_VPD if it exists.
	if err := dutConn.CommandContext(ctx, "vpd", "-i", "RW_VPD", "-g", "ActivateDate").Run(); err == nil {
		if err := dutConn.CommandContext(ctx, "vpd", "-i", "RW_VPD", "-d", "ActivateDate").Run(exec.DumpLogOnError); err != nil {
			return errors.Wrap(err, "failed to delete ActivateDate")
		}
	}

	// Setting the RLZ ping embargo end date to one month ago.
	currentTime := time.Now()
	last1Month := currentTime.AddDate(0, -1, 0)
	timeLayout := "2006-01-02"
	oneMonthAgoDate := last1Month.Format(timeLayout)
	oneMonthAgo := fmt.Sprintf("rlz_embargo_end_date=\"%s\"", oneMonthAgoDate)
	if err := dutConn.CommandContext(ctx, "vpd", "-i", "RW_VPD", "-s", oneMonthAgo).Run(exec.DumpLogOnError); err != nil {
		return errors.Wrap(err, "failed to set rlz date")
	}

	return nil
}

func preProvisionDevice(ctx context.Context, serialNumber, hardwareModel, deviceProvisionToken, customerID, batchKey string) (name string, err error) {
	// Prepare and issue a request.
	bodyCommand := fmt.Sprintf(`{"serialNumber": "%s", "hardwareModel": "%s", "devicePreProvisioningToken": "%s", "attestedDeviceId": "%s", "customerId": "%s"}`, serialNumber, hardwareModel, deviceProvisionToken, serialNumber, customerID)
	urlWithBatchKey := fmt.Sprintf("https://chromecommercial.googleapis.com/v1/preProvisionedDevices?key=%s", batchKey)
	body := strings.NewReader(bodyCommand)
	req, err := http.NewRequest("POST", urlWithBatchKey, body)
	if err != nil {
		return "", errors.Wrap(err, "failed to create request")
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", errors.Wrap(err, "failed to issue request")
	}
	defer resp.Body.Close()

	// Validate response.
	respBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", errors.Wrap(err, "failed to read response")
	}
	type preProvisionedDevice struct {
		Name string
	}
	var parsedResponse preProvisionedDevice
	if err := json.Unmarshal(respBytes, &parsedResponse); err != nil {
		return "", errors.Wrapf(err, "failed to parse response (status code %d): %s", resp.StatusCode, string(respBytes))
	}
	if resp.StatusCode != http.StatusOK || len(parsedResponse.Name) == 0 || !strings.HasPrefix(parsedResponse.Name, "preProvisionedDevices/") {
		return "", errors.Errorf("unsuccessful response (status code %d): %s", resp.StatusCode, string(respBytes))
	}
	testing.ContextLogf(ctx, "Succesfully pre-provisioned device with name %s", parsedResponse.Name)
	return parsedResponse.Name, nil
}

func deletePreProvisioningRecord(ctx context.Context, name, batchKey string) error {
	urlWithBatchKey := fmt.Sprintf("https://chromecommercial.googleapis.com/v1/%s?key=%s", name, batchKey)
	deleteReq, err := http.NewRequest("DELETE", urlWithBatchKey, strings.NewReader(""))
	if err != nil {
		return errors.Wrap(err, "failed to create DELETE request")
	}
	deleteResp, err := http.DefaultClient.Do(deleteReq)
	if err != nil {
		return errors.Wrap(err, "failed to issue DELETE request")
	}
	defer deleteResp.Body.Close()

	if deleteResp.StatusCode != http.StatusOK {
		respBytes, err := io.ReadAll(deleteResp.Body)
		if err != nil {
			return errors.Wrapf(err, "failed to read unsuccessful DELETE response (status code %d)", deleteResp.StatusCode)
		}
		return errors.Errorf("unsuccessful DELETE response (status code %d): %s", deleteResp.StatusCode, string(respBytes))
	}
	testing.ContextLogf(ctx, "Succesfully deleted pre-provisioning record for %s", name)
	return nil
}