"constellation recover" CLI command

Signed-off-by: Malte Poll <mp@edgeless.systems>
2024-10-01 01:36:09 -04:00 · 2022-05-02 13:21:07 +02:00 · 2022-05-02 13:21:07 +02:00 · e13ec3f914
commit e13ec3f914
parent 9d31c9637b
7 changed files with 632 additions and 0 deletions
--- a/cli/cmd/recover.go
+++ b/cli/cmd/recover.go
@ -0,0 +1,155 @@
+package cmd
+
+import (
+	"context"
+	"encoding/base64"
+	"errors"
+	"net"
+	"regexp"
+	"strings"
+
+	"github.com/edgelesssys/constellation/cli/cloud/cloudcmd"
+	"github.com/edgelesssys/constellation/cli/cloudprovider"
+	"github.com/edgelesssys/constellation/cli/file"
+	"github.com/edgelesssys/constellation/cli/proto"
+	"github.com/edgelesssys/constellation/coordinator/util"
+	"github.com/edgelesssys/constellation/internal/config"
+	"github.com/edgelesssys/constellation/internal/constants"
+	"github.com/edgelesssys/constellation/internal/state"
+	"github.com/spf13/afero"
+	"github.com/spf13/cobra"
+)
+
+var diskUUIDRegexp = regexp.MustCompile("^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
+
+func newRecoverCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "recover",
+		Short: "Recover a Constellation.",
+		Long: "Recover a Constellation by sending a recovery key to an instance in the boot stage." +
+			"\nThis is only required if instances restart without other instances available for bootstrapping.",
+		Args: cobra.ExactArgs(0),
+		RunE: runRecover,
+	}
+	cmd.Flags().String("ip", "", "Instance IP address.")
+	_ = cmd.MarkFlagRequired("ip")
+	cmd.Flags().String("disk-uuid", "", "Disk UUID of the encrypted state disk.")
+	_ = cmd.MarkFlagRequired("disk-uuid")
+	cmd.Flags().String("master-secret", "", "Path to base64 encoded master secret.")
+	return cmd
+}
+
+func runRecover(cmd *cobra.Command, args []string) error {
+	fileHandler := file.NewHandler(afero.NewOsFs())
+	recoveryClient := &proto.KeyClient{}
+	defer recoveryClient.Close()
+	return recover(cmd.Context(), cmd, fileHandler, recoveryClient)
+}
+
+func recover(ctx context.Context, cmd *cobra.Command, fileHandler file.Handler, recoveryClient recoveryClient) error {
+	flags, err := parseRecoverFlags(cmd, fileHandler)
+	if err != nil {
+		return err
+	}
+
+	config, err := config.FromFile(fileHandler, flags.devConfigPath)
+	if err != nil {
+		return err
+	}
+
+	var stat state.ConstellationState
+	if err := fileHandler.ReadJSON(constants.StateFilename, &stat); err != nil {
+		return err
+	}
+
+	validators, err := cloudcmd.NewValidators(cloudprovider.FromString(stat.CloudProvider), config)
+	if err != nil {
+		return err
+	}
+	cmd.Print(validators.WarningsIncludeInit())
+
+	if err := recoveryClient.Connect(flags.ip, *config.CoordinatorPort, validators.V()); err != nil {
+		return err
+	}
+
+	diskKey, err := deriveStateDiskKey(flags.masterSecret, flags.diskUUID)
+	if err != nil {
+		return err
+	}
+
+	if err := recoveryClient.PushStateDiskKey(ctx, diskKey); err != nil {
+		return err
+	}
+
+	cmd.Println("Pushed recovery key.")
+	return nil
+}
+
+func parseRecoverFlags(cmd *cobra.Command, fileHandler file.Handler) (recoverFlags, error) {
+	ip, err := cmd.Flags().GetString("ip")
+	if err != nil {
+		return recoverFlags{}, err
+	}
+	if netIP := net.ParseIP(ip); netIP == nil {
+		return recoverFlags{}, errors.New("flag '--ip' isn't a valid IP address")
+	}
+
+	diskUUID, err := cmd.Flags().GetString("disk-uuid")
+	if err != nil {
+		return recoverFlags{}, err
+	}
+	if match := diskUUIDRegexp.MatchString(diskUUID); !match {
+		return recoverFlags{}, errors.New("flag '--disk-uuid' isn't a valid LUKS UUID")
+	}
+	diskUUID = strings.ToLower(diskUUID)
+
+	masterSecretPath, err := cmd.Flags().GetString("master-secret")
+	if err != nil {
+		return recoverFlags{}, err
+	}
+	if masterSecretPath == "" {
+		masterSecretPath = constants.MasterSecretFilename
+	}
+	masterSecret, err := readMasterSecret(fileHandler, masterSecretPath)
+	if err != nil {
+		return recoverFlags{}, err
+	}
+
+	devConfigPath, err := cmd.Flags().GetString("dev-config")
+	if err != nil {
+		return recoverFlags{}, err
+	}
+
+	return recoverFlags{
+		ip:            ip,
+		diskUUID:      diskUUID,
+		masterSecret:  masterSecret,
+		devConfigPath: devConfigPath,
+	}, nil
+}
+
+type recoverFlags struct {
+	ip            string
+	diskUUID      string
+	masterSecret  []byte
+	devConfigPath string
+}
+
+// readMasterSecret reads a base64 encoded master secret from file.
+func readMasterSecret(fileHandler file.Handler, filename string) ([]byte, error) {
+	// Try to read the base64 secret from file
+	encodedSecret, err := fileHandler.Read(filename)
+	if err != nil {
+		return nil, err
+	}
+	decoded, err := base64.StdEncoding.DecodeString(string(encodedSecret))
+	if err != nil {
+		return nil, err
+	}
+	return decoded, nil
+}
+
+// deriveStateDiskKey derives a state disk key from a master secret and a disk UUID.
+func deriveStateDiskKey(masterKey []byte, diskUUID string) ([]byte, error) {
+	return util.DeriveKey(masterKey, []byte("Constellation"), []byte("key"+diskUUID), constants.StateDiskKeyLength)
+}
--- a/cli/cmd/recover_test.go
+++ b/cli/cmd/recover_test.go
@ -0,0 +1,356 @@
+package cmd
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/edgelesssys/constellation/cli/file"
+	"github.com/edgelesssys/constellation/internal/constants"
+	"github.com/edgelesssys/constellation/internal/state"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRecoverCmdArgumentValidation(t *testing.T) {
+	testCases := map[string]struct {
+		args    []string
+		wantErr bool
+	}{
+		"no args":            {[]string{}, false},
+		"too many arguments": {[]string{"abc"}, true},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			cmd := newRecoverCmd()
+			err := cmd.ValidateArgs(tc.args)
+
+			if tc.wantErr {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+			}
+		})
+	}
+}
+
+func TestRecover(t *testing.T) {
+	validState := state.ConstellationState{CloudProvider: "GCP"}
+	invalidCSPState := state.ConstellationState{CloudProvider: "invalid"}
+
+	testCases := map[string]struct {
+		setupFs          func(*require.Assertions) afero.Fs
+		existingState    state.ConstellationState
+		client           *stubRecoveryClient
+		ipFlag           string
+		diskUUIDFlag     string
+		masterSecretFlag string
+		devConfigFlag    string
+		stateless        bool
+		wantErr          bool
+		wantKey          []byte
+	}{
+		"works": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: validState,
+			client:        &stubRecoveryClient{},
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "00000000-0000-0000-0000-000000000000",
+			wantKey:       []byte{0x2e, 0x4d, 0x40, 0x3a, 0x90, 0x96, 0x6e, 0xd, 0x42, 0x3, 0x98, 0xd, 0xce, 0xc5, 0x73, 0x26, 0xf4, 0x87, 0xcf, 0x85, 0x73, 0xe1, 0xb7, 0xd6, 0xb2, 0x82, 0x4c, 0xd9, 0xbc, 0xa5, 0x7c, 0x32},
+		},
+		"uppercase disk uuid works": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: validState,
+			client:        &stubRecoveryClient{},
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "ABCDEFAB-CDEF-ABCD-ABCD-ABCDEFABCDEF",
+			wantKey:       []byte{0xa9, 0x4, 0x3a, 0x74, 0x53, 0xeb, 0x23, 0xb2, 0xbc, 0x88, 0xce, 0xa7, 0x4e, 0xa9, 0xda, 0x9f, 0x11, 0x85, 0xc4, 0x2f, 0x1f, 0x25, 0x10, 0xc9, 0xec, 0xfe, 0xa, 0x6c, 0xa2, 0x6f, 0x53, 0x34},
+		},
+		"lowercase disk uuid results in same key": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: validState,
+			client:        &stubRecoveryClient{},
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "abcdefab-cdef-abcd-abcd-abcdefabcdef",
+			wantKey:       []byte{0xa9, 0x4, 0x3a, 0x74, 0x53, 0xeb, 0x23, 0xb2, 0xbc, 0x88, 0xce, 0xa7, 0x4e, 0xa9, 0xda, 0x9f, 0x11, 0x85, 0xc4, 0x2f, 0x1f, 0x25, 0x10, 0xc9, 0xec, 0xfe, 0xa, 0x6c, 0xa2, 0x6f, 0x53, 0x34},
+		},
+		"missing flags": {
+			setupFs: func(require *require.Assertions) afero.Fs { return afero.NewMemMapFs() },
+			wantErr: true,
+		},
+		"missing dev-config": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "00000000-0000-0000-0000-000000000000",
+			devConfigFlag: "nonexistent-dev-config",
+			wantErr:       true,
+		},
+		"missing state": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: validState,
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "00000000-0000-0000-0000-000000000000",
+			stateless:     true,
+			wantErr:       true,
+		},
+		"invalid cloud provider": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: invalidCSPState,
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "00000000-0000-0000-0000-000000000000",
+			wantErr:       true,
+		},
+		"connect fails": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: validState,
+			client:        &stubRecoveryClient{connectErr: errors.New("connect failed")},
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "00000000-0000-0000-0000-000000000000",
+			wantErr:       true,
+		},
+		"pushing state key fails": {
+			setupFs: func(require *require.Assertions) afero.Fs {
+				fs := afero.NewMemMapFs()
+				require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+				return fs
+			},
+			existingState: validState,
+			client:        &stubRecoveryClient{pushStateDiskKeyErr: errors.New("pushing key failed")},
+			ipFlag:        "192.0.2.1",
+			diskUUIDFlag:  "00000000-0000-0000-0000-000000000000",
+			wantErr:       true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			require := require.New(t)
+
+			cmd := newRecoverCmd()
+			cmd.Flags().String("dev-config", "", "") // register persisten flag manually
+			out := &bytes.Buffer{}
+			cmd.SetOut(out)
+			cmd.SetErr(&bytes.Buffer{})
+			if tc.ipFlag != "" {
+				require.NoError(cmd.Flags().Set("ip", tc.ipFlag))
+			}
+			if tc.diskUUIDFlag != "" {
+				require.NoError(cmd.Flags().Set("disk-uuid", tc.diskUUIDFlag))
+			}
+			if tc.masterSecretFlag != "" {
+				require.NoError(cmd.Flags().Set("master-secret", tc.masterSecretFlag))
+			}
+			if tc.devConfigFlag != "" {
+				require.NoError(cmd.Flags().Set("dev-config", tc.devConfigFlag))
+			}
+			fileHandler := file.NewHandler(tc.setupFs(require))
+			if !tc.stateless {
+				require.NoError(fileHandler.WriteJSON(constants.StateFilename, tc.existingState, file.OptNone))
+			}
+
+			ctx := context.Background()
+			err := recover(ctx, cmd, fileHandler, tc.client)
+
+			if tc.wantErr {
+				assert.Error(err)
+				return
+			}
+
+			assert.NoError(err)
+			assert.Contains(out.String(), "Pushed recovery key.")
+			assert.Equal(tc.wantKey, tc.client.pushStateDiskKeyKey)
+		})
+	}
+}
+
+func TestParseRecoverFlags(t *testing.T) {
+	testCases := map[string]struct {
+		args      []string
+		wantFlags recoverFlags
+		wantErr   bool
+	}{
+		"no flags": {
+			wantErr: true,
+		},
+		"invalid ip": {
+			args:    []string{"--ip", "invalid", "--disk-uuid", "12345678-1234-1234-1234-123456789012"},
+			wantErr: true,
+		},
+		"invalid disk uuid": {
+			args:    []string{"--ip", "192.0.2.1", "--disk-uuid", "invalid"},
+			wantErr: true,
+		},
+		"invalid master secret path": {
+			args:    []string{"--ip", "192.0.2.1", "--disk-uuid", "12345678-1234-1234-1234-123456789012", "--master-secret", "invalid"},
+			wantErr: true,
+		},
+		"minimal args set": {
+			args: []string{"--ip", "192.0.2.1", "--disk-uuid", "12345678-1234-1234-1234-123456789012"},
+			wantFlags: recoverFlags{
+				ip:           "192.0.2.1",
+				diskUUID:     "12345678-1234-1234-1234-123456789012",
+				masterSecret: []byte("constellation-master-secret-leng"),
+			},
+		},
+		"all args set": {
+			args: []string{
+				"--ip", "192.0.2.1", "--disk-uuid", "12345678-1234-1234-1234-123456789012",
+				"--master-secret", "constellation-mastersecret.base64", "--dev-config", "dev-config-path",
+			},
+			wantFlags: recoverFlags{
+				ip:            "192.0.2.1",
+				diskUUID:      "12345678-1234-1234-1234-123456789012",
+				masterSecret:  []byte("constellation-master-secret-leng"),
+				devConfigPath: "dev-config-path",
+			},
+		},
+		"uppercase disk-uuid is converted to lowercase": {
+			args: []string{"--ip", "192.0.2.1", "--disk-uuid", "ABCDEFAB-CDEF-ABCD-ABCD-ABCDEFABCDEF"},
+			wantFlags: recoverFlags{
+				ip:           "192.0.2.1",
+				diskUUID:     "abcdefab-cdef-abcd-abcd-abcdefabcdef",
+				masterSecret: []byte("constellation-master-secret-leng"),
+			},
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			require := require.New(t)
+
+			fs := afero.NewMemMapFs()
+			require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="), 0o777))
+			cmd := newRecoverCmd()
+			cmd.Flags().String("dev-config", "", "") // register persistent flag manually
+			require.NoError(cmd.ParseFlags(tc.args))
+			flags, err := parseRecoverFlags(cmd, file.NewHandler(fs))
+
+			if tc.wantErr {
+				assert.Error(err)
+				return
+			}
+			assert.NoError(err)
+			assert.Equal(tc.wantFlags, flags)
+		})
+	}
+}
+
+func TestReadMasterSecret(t *testing.T) {
+	testCases := map[string]struct {
+		fileContents     []byte
+		filename         string
+		wantMasterSecret []byte
+		wantErr          bool
+	}{
+		"invalid base64": {
+			fileContents: []byte("invalid"),
+			filename:     "constellation-mastersecret.base64",
+			wantErr:      true,
+		},
+		"invalid filename": {
+			fileContents: []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="),
+			filename:     "invalid",
+			wantErr:      true,
+		},
+		"correct master secret": {
+			fileContents:     []byte("Y29uc3RlbGxhdGlvbi1tYXN0ZXItc2VjcmV0LWxlbmc="),
+			filename:         "constellation-mastersecret.base64",
+			wantMasterSecret: []byte("constellation-master-secret-leng"),
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+			require := require.New(t)
+
+			fs := afero.NewMemMapFs()
+			require.NoError(afero.WriteFile(fs, "constellation-mastersecret.base64", tc.fileContents, 0o777))
+			masterSecret, err := readMasterSecret(file.NewHandler(fs), tc.filename)
+
+			if tc.wantErr {
+				assert.Error(err)
+				return
+			}
+			assert.NoError(err)
+			assert.Equal(tc.wantMasterSecret, masterSecret)
+		})
+	}
+}
+
+func TestDeriveStateDiskKey(t *testing.T) {
+	testCases := map[string]struct {
+		masterKey        []byte
+		diskUUID         string
+		wantStateDiskKey []byte
+	}{
+		"all zero": {
+			masterKey: []byte{
+				0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+				0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+			},
+			diskUUID: "00000000-0000-0000-0000-000000000000",
+			wantStateDiskKey: []byte{
+				0xa8, 0xb0, 0x86, 0x83, 0x6f, 0x0b, 0x26, 0x04, 0x86, 0x22, 0x27, 0xcc, 0xa1, 0x1c, 0xaf, 0x6c,
+				0x30, 0x4d, 0x90, 0x89, 0x82, 0x68, 0x53, 0x7f, 0x4f, 0x46, 0x7a, 0x65, 0xa2, 0x5d, 0x5e, 0x43,
+			},
+		},
+		"all 0xff": {
+			masterKey: []byte{
+				0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+				0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			},
+			diskUUID: "ffffffff-ffff-ffff-ffff-ffffffffffff",
+			wantStateDiskKey: []byte{
+				0x24, 0x18, 0x84, 0x7f, 0xca, 0x86, 0x55, 0xb5, 0x45, 0xa6, 0xb3, 0xc4, 0x45, 0xbb, 0x08, 0x10,
+				0x16, 0xb3, 0xde, 0x30, 0x30, 0x74, 0x0b, 0xd4, 0x1e, 0x22, 0x55, 0x45, 0x51, 0x91, 0xfb, 0xa9,
+			},
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			stateDiskKey, err := deriveStateDiskKey(tc.masterKey, tc.diskUUID)
+
+			assert.NoError(err)
+			assert.Equal(tc.wantStateDiskKey, stateDiskKey)
+		})
+	}
+}
--- a/cli/cmd/recoveryclient.go
+++ b/cli/cmd/recoveryclient.go
@ -0,0 +1,14 @@
+package cmd
+
+import (
+	"context"
+	"io"
+
+	"github.com/edgelesssys/constellation/coordinator/atls"
+)
+
+type recoveryClient interface {
+	Connect(ip, port string, validators []atls.Validator) error
+	PushStateDiskKey(ctx context.Context, stateDiskKey []byte) error
+	io.Closer
+}
--- a/cli/cmd/recoveryclient_test.go
+++ b/cli/cmd/recoveryclient_test.go
@ -0,0 +1,31 @@
+package cmd
+
+import (
+	"context"
+
+	"github.com/edgelesssys/constellation/coordinator/atls"
+)
+
+type stubRecoveryClient struct {
+	conn                bool
+	connectErr          error
+	closeErr            error
+	pushStateDiskKeyErr error
+
+	pushStateDiskKeyKey []byte
+}
+
+func (c *stubRecoveryClient) Connect(_, _ string, _ []atls.Validator) error {
+	c.conn = true
+	return c.connectErr
+}
+
+func (c *stubRecoveryClient) Close() error {
+	c.conn = false
+	return c.closeErr
+}
+
+func (c *stubRecoveryClient) PushStateDiskKey(_ context.Context, stateDiskKey []byte) error {
+	c.pushStateDiskKeyKey = stateDiskKey
+	return c.pushStateDiskKeyErr
+}
--- a/cli/cmd/root.go
+++ b/cli/cmd/root.go
@ -58,6 +58,7 @@ func init() {
 	rootCmd.AddCommand(newCreateCmd())
 	rootCmd.AddCommand(newInitCmd())
 	rootCmd.AddCommand(newVerifyCmd())
+	rootCmd.AddCommand(newRecoverCmd())
 	rootCmd.AddCommand(newTerminateCmd())
 	rootCmd.AddCommand(newVersionCmd())
 }
--- a/cli/proto/recover.go
+++ b/cli/proto/recover.go
@ -0,0 +1,70 @@
+package proto
+
+import (
+	"context"
+	"errors"
+	"net"
+
+	"github.com/edgelesssys/constellation/coordinator/atls"
+	"github.com/edgelesssys/constellation/state/keyservice/keyproto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+)
+
+// KeyClient wraps a KeyAPI client and the connection to it.
+type KeyClient struct {
+	conn   *grpc.ClientConn
+	keyapi keyproto.APIClient
+}
+
+// Connect connects the client to a given server, using the handed
+// Validators for the attestation of the connection.
+// The connection must be closed using Close(). If connect is
+// called on a client that already has a connection, the old
+// connection is closed.
+func (c *KeyClient) Connect(ip, port string, validators []atls.Validator) error {
+	addr := net.JoinHostPort(ip, port)
+
+	tlsConfig, err := atls.CreateAttestationClientTLSConfig(validators)
+	if err != nil {
+		return err
+	}
+
+	conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)))
+	if err != nil {
+		return err
+	}
+	_ = c.Close()
+	c.conn = conn
+	c.keyapi = keyproto.NewAPIClient(conn)
+	return nil
+}
+
+// Close closes the grpc connection of the client.
+// Close is idempotent and can be called on non connected clients
+// without returning an error.
+func (c *KeyClient) Close() error {
+	if c.conn == nil {
+		return nil
+	}
+	if err := c.conn.Close(); err != nil {
+		return err
+	}
+	c.conn = nil
+	return nil
+}
+
+// PushStateDiskKey pushes the state disk key to a constellation instance in recovery mode.
+// The state disk key must be derived from the UUID of the state disk and the master key.
+func (c *KeyClient) PushStateDiskKey(ctx context.Context, stateDiskKey []byte) error {
+	if c.keyapi == nil {
+		return errors.New("client is not connected")
+	}
+
+	req := &keyproto.PushStateDiskKeyRequest{
+		StateDiskKey: stateDiskKey,
+	}
+
+	_, err := c.keyapi.PushStateDiskKey(ctx, req)
+	return err
+}
--- a/internal/constants/constants.go
+++ b/internal/constants/constants.go
@ -27,6 +27,11 @@ const (
 	AdminConfFilename     = "constellation-admin.conf"
 	MasterSecretFilename  = "constellation-mastersecret.base64"
 	WGQuickConfigFilename = "wg0.conf"
+
+	//
+	// Cryptographic constants.
+	//
+	StateDiskKeyLength = 32
 )

 // CliVersion is the version of the CLI. Left as a separate variable to allow override during build.