Merge branch 'feature/transfercampaign'

Author: p-b--

cmd/cmd.go

@@ -19,6 +19,7 @@ var (
 	clientKeyFile              string
 	caCertFiles                []string
 	tlsParams                  *util.CustomTlsParams
+	sourceCampaignName         string
 )
 
 // RootCmd is the root command for the CLI

cmd/restorer.go

@@ -111,18 +111,30 @@ var restoreCmd = &cobra.Command{
 		slog.InfoContext(ctx, "validated credentials and fetched vectr version", "hostname", hostname, "vectr-version", vectrVersion)
 		versionContext := context.WithValue(ctx, vat.VECTR_VERSION, vat.VatContextValue(vectrVersion))
 
-		optionalParams := &vat.RestoreOptionalParams{
-			AssessmentName:             targetAssessmentName,
-			OverrideAssessmentTemplate: overrideAssessmentTemplate,
-		}
+		if sourceCampaignName == "" {
+			optionalParams := &vat.RestoreOptionalParams{
+				AssessmentName:             targetAssessmentName,
+				OverrideAssessmentTemplate: overrideAssessmentTemplate,
+			}
 
-		// Restore the assessment
-		if err := vat.RestoreAssessment(versionContext, client, db, &assessmentData, optionalParams); err != nil {
-			slog.ErrorContext(versionContext, "Failed to restore assessment", "error", err)
-			os.Exit(1)
+			// Restore the assessment
+			if err := vat.RestoreAssessment(versionContext, client, db, &assessmentData, optionalParams); err != nil {
+				slog.ErrorContext(versionContext, "Failed to restore assessment", "error", err)
+				os.Exit(1)
+			}
+			slog.InfoContext(ctx, "Assessment restored successfully")
+		} else {
+			if targetAssessmentName == "" {
+				slog.ErrorContext(ctx, "--target-assessment-name is required when using --source-campaign-name")
+				os.Exit(1)
+			}
+			slog.InfoContext(ctx, "Restoring campaign", "source-campaign", sourceCampaignName, "target-assessment", targetAssessmentName)
+			if err := vat.RestoreCampaign(versionContext, client, db, &assessmentData, sourceCampaignName, targetAssessmentName); err != nil {
+				slog.ErrorContext(versionContext, "Failed to restore campaign", "error", err)
+				os.Exit(1)
+			}
+			slog.InfoContext(ctx, "Campaign restored successfully")
 		}
-
-		slog.InfoContext(ctx, "Assessment restored successfully")
 	},
 }
 
@@ -134,8 +146,9 @@ func init() {
 	restoreCmd.Flags().StringVar(&credentialsFile, "vectr-creds-file", "", "Path to the credentials file (required)")
 	restoreCmd.Flags().StringVar(&inputFile, "input-file", "", "Path to the encrypted input file (required)")
 	restoreCmd.Flags().StringVar(&passphraseFile, "passphrase-file", "", "Path to the file containing the decryption passphrase")
-	restoreCmd.Flags().StringVar(&targetAssessmentName, "target-assessment-name", "", "The assessment name to set in the new instance")
+	restoreCmd.Flags().StringVar(&targetAssessmentName, "target-assessment-name", "", "The assessment name to set in the new instance. Required when using --source-campaign-name.")
 	restoreCmd.Flags().BoolVar(&overrideAssessmentTemplate, "override-template-assessment", false, "Override any set template name in the serialized data and load template test cases anyway")
+	restoreCmd.Flags().StringVar(&sourceCampaignName, "source-campaign-name", "", "Name of a specific campaign to restore from the input file. If set, --target-assessment-name must be an existing assessment.")
 
 	// Mark flags as required
 	restoreCmd.MarkFlagsOneRequired("db", "env")

cmd/transfer.go

@@ -100,15 +100,28 @@ var transferCmd = &cobra.Command{
 			os.Exit(1)
 		}
 
-		optionalParams := &vat.RestoreOptionalParams{
-			AssessmentName:             targetAssessmentName,
-			OverrideAssessmentTemplate: overrideAssessmentTemplate,
-		}
-		// Transfer the assessment data to the target instance
-		slog.InfoContext(targetVersionContext, "Transferring assessment data to target instance", "hostname", targetHostname, "db", targetDB)
-		if err := vat.RestoreAssessment(targetVersionContext, targetClient, targetDB, assessmentData, optionalParams); err != nil {
-			slog.ErrorContext(targetVersionContext, "Failed to transfer assessment data to target instance", "error", err)
-			os.Exit(1)
+		if sourceCampaignName == "" {
+			optionalParams := &vat.RestoreOptionalParams{
+				AssessmentName:             targetAssessmentName,
+				OverrideAssessmentTemplate: overrideAssessmentTemplate,
+			}
+			// Original full assessment transfer logic
+			slog.InfoContext(targetVersionContext, "Transferring assessment data to target instance", "hostname", targetHostname, "db", targetDB)
+			if err := vat.RestoreAssessment(targetVersionContext, targetClient, targetDB, assessmentData, optionalParams); err != nil {
+				slog.ErrorContext(targetVersionContext, "Failed to transfer assessment data to target instance", "error", err)
+				os.Exit(1)
+			}
+		} else {
+			// New campaign-only transfer logic
+			if targetAssessmentName == "" {
+				slog.ErrorContext(ctx, "--target-assessment-name is required when using --source-campaign-name")
+				os.Exit(1)
+			}
+			slog.InfoContext(targetVersionContext, "Transferring campaign to target assessment", "source-campaign", sourceCampaignName, "target-assessment", targetAssessmentName)
+			if err := vat.RestoreCampaign(targetVersionContext, targetClient, targetDB, assessmentData, sourceCampaignName, targetAssessmentName); err != nil {
+				slog.ErrorContext(targetVersionContext, "Failed to transfer campaign to target instance", "error", err)
+				os.Exit(1)
+			}
 		}
 
 		slog.InfoContext(ctx, "Assessment transferred successfully")
@@ -128,6 +141,7 @@ func init() {
 	transferCmd.Flags().StringVar(&assessmentName, "assessment-name", "", "Name of the assessment to transfer (required)")
 	transferCmd.Flags().StringVar(&targetAssessmentName, "target-assessment-name", "", "The assessment name to set in the new instance")
 	transferCmd.Flags().BoolVar(&overrideAssessmentTemplate, "override-template-assessment", false, "Ignore the template name in the serialized data and load template test cases anyway")
+	transferCmd.Flags().StringVar(&sourceCampaignName, "source-campaign-name", "", "Name of a specific campaign to transfer. If set, --target-assessment-name must be an existing assessment.")
 
 	// Mark flags as required
 	transferCmd.MarkFlagRequired("source-hostname")