From 64477712c59076ab83ec0ba09cd48789190ef21d Mon Sep 17 00:00:00 2001
From: Keepers <104464746+ryanfkeepers@users.noreply.github.com>
Date: Thu, 16 Jun 2022 17:24:10 -0600
Subject: [PATCH] add flags to restore exchange (#213)

* add flags to restore exchange

Adds the user, mail, and folder flags to corso restore exchange.
---
 src/cli/restore/exchange.go      | 35 ++++++++++++++++++++++++----
 src/cli/restore/exchange_test.go | 40 ++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 5 deletions(-)
 create mode 100644 src/cli/restore/exchange_test.go

diff --git a/src/cli/restore/exchange.go b/src/cli/restore/exchange.go
index dc775aebf..89cf44b18 100644
--- a/src/cli/restore/exchange.go
+++ b/src/cli/restore/exchange.go
@@ -14,20 +14,25 @@ import (
 
 // exchange bucket info from flags
 var (
-	user string
+	folder string
+	mail   string
+	user   string
 )
 
 // called by restore.go to map parent subcommands to provider-specific handling.
 func addExchangeApp(parent *cobra.Command) *cobra.Command {
-	parent.AddCommand(exchangeCreateCmd)
+	parent.AddCommand(exchangeCmd)
 
-	// todo (keepers): add flags.
+	fs := exchangeCmd.Flags()
+	fs.StringVar(&user, "user", "", "ID of the user whose echange data will get restored.")
+	fs.StringVar(&folder, "folder", "", "Name of the mail folder being restored.")
+	fs.StringVar(&mail, "mail", "", "ID of the mail message being restored.")
 
-	return exchangeCreateCmd
+	return exchangeCmd
 }
 
 // `corso restore create exchange [<flag>...]`
-var exchangeCreateCmd = &cobra.Command{
+var exchangeCmd = &cobra.Command{
 	Use:   "exchange",
 	Short: "Restore M365 Exchange service data",
 	RunE:  createExchangeCmd,
@@ -36,6 +41,15 @@ var exchangeCreateCmd = &cobra.Command{
 
 // initializes a s3 repo.
 func createExchangeCmd(cmd *cobra.Command, args []string) error {
+	if cmd.Flags().NFlag() == 0 {
+		cmd.Help()
+		return nil
+	}
+
+	if err := validateRestoreFlags(user, folder, mail); err != nil {
+		return errors.Wrap(err, "Missing required flags")
+	}
+
 	s, cfgTenantID, err := config.MakeS3Config(true, nil)
 	if err != nil {
 		return err
@@ -69,3 +83,14 @@ func createExchangeCmd(cmd *cobra.Command, args []string) error {
 	fmt.Printf("Restored Exchange in %s for user %s.\n", s.Provider, user)
 	return nil
 }
+
+func validateRestoreFlags(u, f, m string) error {
+	lu, lf, lm := len(u), len(f), len(m)
+	if (lu == 0 || u == "*") && (lf+lm > 0) {
+		return errors.New("a specific --user must be provided if --folder or --mail is specified")
+	}
+	if (lf == 0 || f == "*") && lm > 0 {
+		return errors.New("a specific --folder must be provided if a --mail is specified")
+	}
+	return nil
+}
diff --git a/src/cli/restore/exchange_test.go b/src/cli/restore/exchange_test.go
new file mode 100644
index 000000000..2629e8f0f
--- /dev/null
+++ b/src/cli/restore/exchange_test.go
@@ -0,0 +1,40 @@
+package restore
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/suite"
+)
+
+type RestoreSuite struct {
+	suite.Suite
+}
+
+func TestRestoreSuite(t *testing.T) {
+	suite.Run(t, new(RestoreSuite))
+}
+
+func (suite *RestoreSuite) TestValidateRestoreFlags() {
+	table := []struct {
+		name     string
+		u, f, m  string
+		errCheck assert.ErrorAssertionFunc
+	}{
+		{"all populated", "u", "f", "m", assert.NoError},
+		{"folder missing user", "", "f", "m", assert.Error},
+		{"folder with wildcard user", "*", "f", "m", assert.Error},
+		{"mail missing user", "", "", "m", assert.Error},
+		{"mail missing folder", "u", "", "m", assert.Error},
+		{"mail with wildcard folder", "u", "*", "m", assert.Error},
+		{"all missing", "", "", "", assert.NoError},
+	}
+	for _, test := range table {
+		suite.T().Run(test.name, func(t *testing.T) {
+			test.errCheck(
+				t,
+				validateRestoreFlags(test.u, test.f, test.m),
+			)
+		})
+	}
+}