restore.go

/*
Copyright AppsCode Inc. and Contributors

Licensed under the AppsCode Free Trial License 1.0.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://github.com/appscode/licenses/raw/1.0.0/AppsCode-Free-Trial-1.0.0.md

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package pkg

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	api_v1beta1 "stash.appscode.dev/apimachinery/apis/stash/v1beta1"
	stash_cs "stash.appscode.dev/apimachinery/client/clientset/versioned"
	stash_cs_util "stash.appscode.dev/apimachinery/client/clientset/versioned/typed/stash/v1beta1/util"
	"stash.appscode.dev/apimachinery/pkg/restic"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	license "go.bytebuilders.dev/license-verifier/kubernetes"
	"gomodules.xyz/flags"
	"gomodules.xyz/go-sh"
	"gomodules.xyz/pointer"
	core "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog/v2"
	appcatalog "kmodules.xyz/custom-resources/apis/appcatalog/v1alpha1"
	appcatalog_cs "kmodules.xyz/custom-resources/client/clientset/versioned"
	v1 "kmodules.xyz/offshoot-api/api/v1"
	"kubedb.dev/apimachinery/apis/config/v1alpha1"
)

var cleanupRestoreFuncs []func() error

func NewCmdRestore() *cobra.Command {
	var (
		masterURL      string
		kubeconfigPath string
		opt            = mongoOptions{
			waitTimeout: 300,
			setupOptions: restic.SetupOptions{
				ScratchDir:  restic.DefaultScratchDir,
				EnableCache: false,
			},
			defaultDumpOptions: restic.DumpOptions{
				Host:     restic.DefaultHost,
				FileName: MongoDumpFile,
			},
		}
	)

	cmd := &cobra.Command{
		Use:               "restore-mongo",
		Short:             "Restores MongoDB Backup",
		DisableAutoGenTag: true,
		PreRunE: func(cmd *cobra.Command, args []string) error {
			return checkCommandExists()
		},
		RunE: func(cmd *cobra.Command, args []string) error {
			defer cleanupRestore()
			flags.EnsureRequiredFlags(cmd, "appbinding", "provider", "storage-secret-name", "storage-secret-namespace")

			// prepare client
			config, err := clientcmd.BuildConfigFromFlags(masterURL, kubeconfigPath)
			if err != nil {
				return err
			}
			opt.config = config

			opt.stashClient, err = stash_cs.NewForConfig(config)
			if err != nil {
				return err
			}
			opt.kubeClient, err = kubernetes.NewForConfig(config)
			if err != nil {
				return err
			}
			opt.catalogClient, err = appcatalog_cs.NewForConfig(config)
			if err != nil {
				return err
			}

			targetRef := api_v1beta1.TargetRef{
				APIVersion: appcatalog.SchemeGroupVersion.String(),
				Kind:       appcatalog.ResourceKindApp,
				Name:       opt.appBindingName,
				Namespace:  opt.appBindingNamespace,
			}
			var restoreOutput *restic.RestoreOutput
			restoreOutput, err = opt.restoreMongoDB(targetRef)
			if err != nil {
				restoreOutput = &restic.RestoreOutput{
					RestoreTargetStatus: api_v1beta1.RestoreMemberStatus{
						Ref:   targetRef,
						Stats: opt.getHostRestoreStats(err),
					},
				}
			}
			// If output directory specified, then write the output in "output.json" file in the specified directory
			if opt.outputDir != "" {
				return restoreOutput.WriteOutput(filepath.Join(opt.outputDir, restic.DefaultOutputFileName))
			}

			return nil
		},
	}

	cmd.Flags().StringVar(&opt.mongoArgs, "mongo-args", opt.mongoArgs, "Additional arguments")
	cmd.Flags().Int32Var(&opt.waitTimeout, "wait-timeout", opt.waitTimeout, "Number of seconds to wait for the database to be ready")

	cmd.Flags().StringVar(&masterURL, "master", masterURL, "The address of the Kubernetes API server (overrides any value in kubeconfig)")
	cmd.Flags().StringVar(&kubeconfigPath, "kubeconfig", kubeconfigPath, "Path to kubeconfig file with authorization information (the master location is set by the master flag).")
	cmd.Flags().StringVar(&opt.namespace, "namespace", "default", "Namespace of Backup/Restore Session")
	cmd.Flags().StringVar(&opt.appBindingName, "appbinding", opt.appBindingName, "Name of the app binding")
	cmd.Flags().StringVar(&opt.appBindingNamespace, "appbinding-namespace", opt.appBindingNamespace, "Namespace of the app binding")
	cmd.Flags().StringVar(&opt.restoreSessionName, "restoresession", opt.restoreSessionName, "Name of the respective RestoreSession object")
	cmd.Flags().IntVar(&opt.maxConcurrency, "max-concurrency", 3, "maximum concurrent backup process to run to take backup from each replicasets")

	cmd.Flags().StringVar(&opt.setupOptions.Provider, "provider", opt.setupOptions.Provider, "Backend provider (i.e. gcs, s3, azure etc)")
	cmd.Flags().StringVar(&opt.setupOptions.Bucket, "bucket", opt.setupOptions.Bucket, "Name of the cloud bucket/container (keep empty for local backend)")
	cmd.Flags().StringVar(&opt.setupOptions.Endpoint, "endpoint", opt.setupOptions.Endpoint, "Endpoint for s3/s3 compatible backend or REST server URL")
	cmd.Flags().StringVar(&opt.setupOptions.Region, "region", opt.setupOptions.Region, "Region for s3/s3 compatible backend")
	cmd.Flags().StringVar(&opt.setupOptions.Path, "path", opt.setupOptions.Path, "Directory inside the bucket where backup will be stored")
	cmd.Flags().StringVar(&opt.setupOptions.ScratchDir, "scratch-dir", opt.setupOptions.ScratchDir, "Temporary directory")
	cmd.Flags().BoolVar(&opt.setupOptions.EnableCache, "enable-cache", opt.setupOptions.EnableCache, "Specify whether to enable caching for restic")
	cmd.Flags().Int64Var(&opt.setupOptions.MaxConnections, "max-connections", opt.setupOptions.MaxConnections, "Specify maximum concurrent connections for GCS, Azure and B2 backend")

	cmd.Flags().StringVar(&opt.storageSecret.Name, "storage-secret-name", opt.storageSecret.Name, "Name of the storage secret")
	cmd.Flags().StringVar(&opt.storageSecret.Namespace, "storage-secret-namespace", opt.storageSecret.Namespace, "Namespace of the storage secret")
	cmd.Flags().StringVar(&opt.authenticationDatabase, "authentication-database", "admin", "Specify the authentication database")

	cmd.Flags().StringVar(&opt.defaultDumpOptions.Host, "hostname", opt.defaultDumpOptions.Host, "Name of the host machine")
	cmd.Flags().StringVar(&opt.defaultDumpOptions.SourceHost, "source-hostname", opt.defaultDumpOptions.SourceHost, "Name of the host whose data will be restored")
	cmd.Flags().StringVar(&opt.defaultDumpOptions.Snapshot, "snapshot", opt.defaultDumpOptions.Snapshot, "Snapshot to dump")

	cmd.Flags().StringVar(&opt.outputDir, "output-dir", opt.outputDir, "Directory where output.json file will be written (keep empty if you don't need to write output in file)")

	return cmd
}

func cleanupRestore() {
	for _, f := range cleanupRestoreFuncs {
		if err := f(); err != nil {
			klog.Errorln(err)
		}
	}
}

func (opt *mongoOptions) restoreMongoDB(targetRef api_v1beta1.TargetRef) (*restic.RestoreOutput, error) {
	var err error
	err = license.CheckLicenseEndpoint(opt.config, licenseApiService, SupportedProducts)
	if err != nil {
		return nil, err
	}

	opt.setupOptions.StorageSecret, err = opt.kubeClient.CoreV1().Secrets(opt.storageSecret.Namespace).Get(context.TODO(), opt.storageSecret.Name, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	// apply nice, ionice settings from env
	opt.setupOptions.Nice, err = v1.NiceSettingsFromEnv()
	if err != nil {
		return nil, err
	}
	opt.setupOptions.IONice, err = v1.IONiceSettingsFromEnv()
	if err != nil {
		return nil, err
	}

	appBinding, err := opt.catalogClient.AppcatalogV1alpha1().AppBindings(opt.appBindingNamespace).Get(context.TODO(), opt.appBindingName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	authSecret, err := opt.kubeClient.CoreV1().Secrets(opt.appBindingNamespace).Get(context.TODO(), appBinding.Spec.Secret.Name, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	err = appBinding.TransformSecret(opt.kubeClient, authSecret.Data)
	if err != nil {
		return nil, err
	}

	var tlsSecret *core.Secret
	if appBinding.Spec.TLSSecret != nil {
		tlsSecret, err = opt.kubeClient.CoreV1().Secrets(opt.appBindingNamespace).Get(context.TODO(), appBinding.Spec.TLSSecret.Name, metav1.GetOptions{})
		if err != nil {
			return nil, err
		}
	}

	hostname, err := appBinding.Hostname()
	if err != nil {
		return nil, err
	}

	port, err := appBinding.Port()
	if err != nil {
		return nil, err
	}

	// unmarshal parameter is the field has value
	parameters := v1alpha1.MongoDBConfiguration{}
	if appBinding.Spec.Parameters != nil {
		if err = json.Unmarshal(appBinding.Spec.Parameters.Raw, &parameters); err != nil {
			klog.Errorf("unable to unmarshal appBinding.Spec.Parameters.Raw. Reason: %v", err)
		}
	}

	// Stash operator does not know how many hosts this plugin will restore. It sets totalHosts field of respective RestoreSession to 1.
	// We must update the totalHosts field to the actual number of hosts it will restore.
	// Otherwise, RestoreSession will stuck in "Running" state.
	// Total hosts for MongoDB:
	// 1. For stand-alone MongoDB, totalHosts=1.
	// 2. For MongoDB ReplicaSet, totalHosts=1.
	// 3. For sharded MongoDB, totalHosts=(number of shard + 1) // extra 1 for config server
	// So, for stand-alone MongoDB and MongoDB ReplicaSet, we don't have to do anything.
	// We only need to update totalHosts field for sharded MongoDB

	opt.totalHosts = 1
	// For sharded MongoDB, parameter.ConfigServer will not be empty
	if parameters.ConfigServer != "" {
		opt.totalHosts = len(parameters.ReplicaSets) + 1 // for each shard there will be one key in parameters.ReplicaSet
		restoreSession, err := opt.stashClient.StashV1beta1().RestoreSessions(opt.namespace).Get(context.TODO(), opt.restoreSessionName, metav1.GetOptions{})
		if err != nil {
			return nil, err
		}
		_, err = stash_cs_util.UpdateRestoreSessionStatus(
			context.TODO(),
			opt.stashClient.StashV1beta1(),
			restoreSession.ObjectMeta,
			func(status *api_v1beta1.RestoreSessionStatus) (types.UID, *api_v1beta1.RestoreSessionStatus) {
				status.TotalHosts = pointer.Int32P(int32(opt.totalHosts))
				return restoreSession.UID, status
			},
			metav1.UpdateOptions{},
		)
		if err != nil {
			return nil, err
		}
	}

	if appBinding.Spec.ClientConfig.CABundle != nil {
		if err := os.WriteFile(filepath.Join(opt.setupOptions.ScratchDir, MongoTLSCertFileName), appBinding.Spec.ClientConfig.CABundle, os.ModePerm); err != nil {
			return nil, errors.Wrap(err, "failed to write key for CA certificate")
		}
		mongoCreds = []interface{}{
			"--tls",
			"--tlsCAFile", filepath.Join(opt.setupOptions.ScratchDir, MongoTLSCertFileName),
			"--tlsCertificateKeyFile", filepath.Join(opt.setupOptions.ScratchDir, MongoClientPemFileName),
		}
		dumpCreds = []interface{}{
			"--ssl",
			"--sslCAFile", filepath.Join(opt.setupOptions.ScratchDir, MongoTLSCertFileName),
			"--sslPEMKeyFile", filepath.Join(opt.setupOptions.ScratchDir, MongoClientPemFileName),
		}

		// get certificate secret to get client certificate
		var pemBytes []byte
		var ok bool
		pemBytes, ok = tlsSecret.Data[MongoClientPemFileName]
		if !ok {
			crt, ok := tlsSecret.Data[core.TLSCertKey]
			if !ok {
				return nil, errors.Wrap(err, "unable to retrieve tls.crt from secret.")
			}
			key, ok := tlsSecret.Data[core.TLSPrivateKeyKey]
			if !ok {
				return nil, errors.Wrap(err, "unable to retrieve tls.key from secret.")
			}
			pemBytes = append(crt[:], []byte("\n")...)
			pemBytes = append(pemBytes, key...)
		}
		if err := os.WriteFile(filepath.Join(opt.setupOptions.ScratchDir, MongoClientPemFileName), pemBytes, os.ModePerm); err != nil {
			return nil, errors.Wrap(err, "failed to write client certificate")
		}
		user, err := getSSLUser(filepath.Join(opt.setupOptions.ScratchDir, MongoClientPemFileName))
		if err != nil {
			return nil, errors.Wrap(err, "unable to get user from ssl.")
		}
		userAuth := []interface{}{
			"-u", user,
			"--authenticationMechanism", "MONGODB-X509",
			"--authenticationDatabase", "$external",
		}
		mongoCreds = append(mongoCreds, userAuth...)

	} else {
		userAuth := []interface{}{
			fmt.Sprintf("--username=%s", authSecret.Data[MongoUserKey]),
			fmt.Sprintf("--password=%s", authSecret.Data[MongoPasswordKey]),
			"--authenticationDatabase", opt.authenticationDatabase,
		}
		mongoCreds = append(mongoCreds, userAuth...)
	}

	err = opt.workOnSuperUser(parameters, hostname)
	if err != nil {
		return nil, err
	}

	getDumpOpts := func(mongoDSN, hostKey string, isStandalone bool) restic.DumpOptions {
		klog.Infoln("processing backupOptions for ", mongoDSN)
		dumpOpt := restic.DumpOptions{
			Host:       hostKey,
			SourceHost: hostKey,
			FileName:   opt.defaultDumpOptions.FileName,
			Snapshot:   opt.defaultDumpOptions.Snapshot,
		}

		// setup pipe command
		restoreCmd := restic.Command{
			Name: MongoRestoreCMD,
			Args: append([]interface{}{
				"--host", mongoDSN,
				"--archive",
			}, dumpCreds...),
		}

		userArgs := strings.Fields(opt.mongoArgs)
		if isStandalone {
			restoreCmd.Args = append(restoreCmd.Args, fmt.Sprintf("--port=%d", port))
		} else {
			// - port is already added in mongoDSN with replicasetName/host:port format.
			// - oplog is enabled automatically for replicasets.
			// Don't use --oplogReplay if user specify any of these arguments through opt.mongoArgs
			// 1. --db
			// 2. --collection
			// 3. --nsInclude
			// 4. --nsExclude
			// xref: https://docs.mongodb.com/manual/reference/program/mongorestore/#cmdoption-mongorestore-oplogreplay
			forbiddenArgs := sets.NewString(
				"-d", "--db",
				"-c", "--collection",
				"--nsInclude",
				"--nsExclude",
			)
			if !containsArg(userArgs, forbiddenArgs) {
				restoreCmd.Args = append(restoreCmd.Args, "--oplogReplay")
			}
		}

		for _, arg := range userArgs {
			restoreCmd.Args = append(restoreCmd.Args, arg)
		}

		// add the restore command to the pipeline
		dumpOpt.StdoutPipeCommands = append(dumpOpt.StdoutPipeCommands, restoreCmd)
		return dumpOpt
	}

	// set opt.maxConcurrency
	if len(parameters.ReplicaSets) <= 1 {
		opt.maxConcurrency = 1
	}

	// If parameters.ReplicaSets is not empty, then replicaset hosts are given in key:value pair,
	// where, keys are host-0,host-1 etc and values are the replicaset dsn of one replicaset component
	//
	// Procedure of restore in a sharded or replicaset cluster
	// - Restore the CSRS primary mongod data files.
	// - Restore Each Shard Replica Set
	// ref: https://docs.mongodb.com/manual/tutorial/backup-sharded-cluster-with-database-dumps/

	if parameters.ConfigServer != "" {
		opt.dumpOptions = append(opt.dumpOptions, getDumpOpts(parameters.ConfigServer, MongoConfigSVRHostKey, false))
	}

	for key, host := range parameters.ReplicaSets {
		opt.dumpOptions = append(opt.dumpOptions, getDumpOpts(host, key, false))
	}

	// if parameters.ReplicaSets is nil, then perform normal backup with clientconfig.Service.Name mongo dsn
	if parameters.ReplicaSets == nil {
		opt.dumpOptions = append(opt.dumpOptions, getDumpOpts(hostname, restic.DefaultHost, true))
	}

	klog.Infoln("processing restore.")

	waitForDBReady(hostname, port, opt.waitTimeout)

	resticWrapper, err := restic.NewResticWrapper(opt.setupOptions)
	if err != nil {
		return nil, err
	}
	// hide password, don't print cmd
	resticWrapper.HideCMD()

	klog.Infof("==> %+v \n %+v \n", opt.dumpOptions, targetRef)

	// Run dump
	out, err := resticWrapper.ParallelDump(opt.dumpOptions, targetRef, opt.maxConcurrency)
	if err != nil {
		klog.Warningln("restore failed!", err.Error())
	}
	// error not returned, error is encoded into output
	return out, nil
}

const (
	SuperRoleName = "interalUseOnlyOplogRestore"
	SuperUserName = "superoplogger"
)

func (opt *mongoOptions) workOnSuperUser(parameters v1alpha1.MongoDBConfiguration, hostNameForStandalone string) error {
	pass, err := genPassword(12)
	if err != nil {
		return err
	}

	if parameters.ConfigServer != "" {
		dsn := parameters.ConfigServer
		err := createSuperRoleAndUser(dsn, pass)
		if err != nil {
			return err
		}
	}
	for key := range parameters.ReplicaSets {
		dsn := parameters.ReplicaSets[key]
		err := createSuperRoleAndUser(dsn, pass)
		if err != nil {
			return err
		}
	}

	if parameters.ReplicaSets == nil {
		err := createSuperRoleAndUser(hostNameForStandalone, pass)
		if err != nil {
			return err
		}
	}

	superUserAuth := []interface{}{
		fmt.Sprintf("--username=%s", SuperUserName),
		fmt.Sprintf("--password=%s", pass),
		"--authenticationDatabase", opt.authenticationDatabase,
	}
	dumpCreds = append(dumpCreds, superUserAuth...)
	return nil
}

func createSuperRoleAndUser(mongoDSN, pass string) error {
	klog.Infof("creating SuperRole & User for %s \n", mongoDSN)
	err := createSuperRole(mongoDSN)
	if err != nil {
		return err
	}

	err = createSuperUser(mongoDSN, pass)
	cleanupRestoreFuncs = append(cleanupRestoreFuncs, func() error {
		return deleteSuperUser(mongoDSN)
	})
	return err
}

func createSuperRole(mongoDSN string) error {
	exists, err := checkRoleExists(mongoDSN, SuperRoleName)
	if err != nil {
		return err
	}
	if !exists {
		klog.Infoln("creating role " + SuperRoleName)
		v := make(map[string]interface{})

		args := append([]interface{}{
			"admin",
			"--host", mongoDSN,
			"--quiet",
			"--eval", `JSON.stringify(db.runCommand({createRole: "` + SuperRoleName + `",privileges:[{resource:{anyResource:true},actions:["anyAction"]}],roles: []}))`,
		}, mongoCreds...)

		if err := sh.Command(MongoCMD, args...).Command("/usr/bin/tail", "-1").UnmarshalJSON(&v); err != nil {
			return err
		}

		if val, ok := v["ok"].(float64); !ok || int(val) != 1 {
			return fmt.Errorf("unable to create role %v. got response: %v", SuperRoleName, v)
		}
	}
	return nil
}

func createSuperUser(mongoDSN, pass string) error {
	exists, err := checkUserExists(mongoDSN, SuperUserName)
	if err != nil {
		return err
	}
	if !exists {
		klog.Infoln("creating user " + SuperUserName)
		v := make(map[string]interface{})

		args := append([]interface{}{
			"admin",
			"--host", mongoDSN,
			"--quiet",
			"--eval", `JSON.stringify(db.runCommand({createUser: "` + SuperUserName + `" ,pwd: "` + pass + `", roles:[{role:"root", db:"admin"},"__system","interalUseOnlyOplogRestore","backup"]}))`,
		}, mongoCreds...)
		if err := sh.Command(MongoCMD, args...).Command("/usr/bin/tail", "-1").UnmarshalJSON(&v); err != nil {
			return err
		}

		if val, ok := v["ok"].(float64); !ok || int(val) != 1 {
			return fmt.Errorf("unable to create user %v. got response: %v", SuperUserName, v)
		}
	}
	return nil
}

func deleteSuperUser(mongoDSN string) error {
	exists, err := checkUserExists(mongoDSN, SuperUserName)
	if err != nil {
		return err
	}
	if exists {
		klog.Infoln("deleting user " + SuperUserName)
		v := make(map[string]interface{})

		args := append([]interface{}{
			"admin",
			"--host", mongoDSN,
			"--quiet",
			"--eval", `JSON.stringify(db.runCommand({dropUser: "` + SuperUserName + `"}))`,
		}, mongoCreds...)
		if err := sh.Command(MongoCMD, args...).Command("/usr/bin/tail", "-1").UnmarshalJSON(&v); err != nil {
			return err
		}

		if val, ok := v["ok"].(float64); !ok || int(val) != 1 {
			return fmt.Errorf("unable to delete user %v. got response: %v", SuperUserName, v)
		}
	}
	return nil
}

func (opt *mongoOptions) getHostRestoreStats(err error) []api_v1beta1.HostRestoreStats {
	var restoreStats []api_v1beta1.HostRestoreStats

	errMsg := fmt.Sprintf("failed to start data restoration: %s", err.Error())
	for _, dumpOpt := range opt.dumpOptions {
		restoreStats = append(restoreStats, api_v1beta1.HostRestoreStats{
			Hostname: dumpOpt.Host,
			Phase:    api_v1beta1.HostRestoreFailed,
			Error:    errMsg,
		})
	}

	if opt.totalHosts > len(restoreStats) {
		rem := opt.totalHosts - len(restoreStats)
		for i := 0; i < rem; i++ {
			restoreStats = append(restoreStats, api_v1beta1.HostRestoreStats{
				Hostname: fmt.Sprintf("unknown-%s", strconv.Itoa(i)),
				Phase:    api_v1beta1.HostRestoreFailed,
				Error:    errMsg,
			})
		}
	}

	return restoreStats
}