From 5ad594299ae04e3a169e8567239371a966fc564e Mon Sep 17 00:00:00 2001
From: Graham Davison <gdavison@hashicorp.com>
Date: Fri, 31 Jul 2020 10:58:47 -0700
Subject: [PATCH] Adds zero-termination option for the `blocks` command

---
 cli/blocks_test.go | 41 +++++++++++++++++++++++++++++++++++++++--
 cli/cmds.go        | 19 ++++++++++++++-----
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/cli/blocks_test.go b/cli/blocks_test.go
index 0a4801a1..0d4761b9 100644
--- a/cli/blocks_test.go
+++ b/cli/blocks_test.go
@@ -188,7 +188,7 @@ func TestCmdBlocksDefault(t *testing.T) {
 			var outB strings.Builder
 			var errB strings.Builder
 			log := common.CreateLogger(&errB)
-			err := findBlocksInFile(fs, log, testcase.sourcefile, false, nil, &outB, &errB)
+			err := findBlocksInFile(fs, log, testcase.sourcefile, false, false, nil, &outB, &errB)
 			actualStdOut := outB.String()
 			actualStdErr := errB.String()
 
@@ -220,7 +220,7 @@ func TestCmdBlocksVerbose(t *testing.T) {
 			var outB strings.Builder
 			var errB strings.Builder
 			log := common.CreateLogger(&errB)
-			err := findBlocksInFile(fs, log, testcase.sourcefile, true, nil, &outB, &errB)
+			err := findBlocksInFile(fs, log, testcase.sourcefile, true, false, nil, &outB, &errB)
 			actualStdErr := errB.String()
 			if err != nil {
 				t.Fatalf("Case %q: Got an error when none was expected: %v", testcase.name, err)
@@ -235,3 +235,40 @@ func TestCmdBlocksVerbose(t *testing.T) {
 		})
 	}
 }
+
+func TestCmdBlocksZeroTerminated(t *testing.T) {
+	t.Parallel()
+
+	for _, testcase := range testcases {
+		t.Run(testcase.name, func(t *testing.T) {
+			t.Parallel()
+
+			fs := afero.NewReadOnlyFs(afero.NewOsFs())
+
+			expectedBuilder := strings.Builder{}
+			for _, block := range testcase.expectedBlocks {
+				fmt.Fprint(&expectedBuilder, block.text, "\n\x00")
+			}
+			expected := expectedBuilder.String()
+
+			var outB strings.Builder
+			var errB strings.Builder
+			log := common.CreateLogger(&errB)
+			err := findBlocksInFile(fs, log, testcase.sourcefile, false, true, nil, &outB, &errB)
+			actualStdOut := outB.String()
+			actualStdErr := errB.String()
+
+			if err != nil {
+				t.Fatalf("Case %q: Got an error when none was expected: %v", testcase.name, err)
+			}
+
+			if actualStdOut != expected {
+				t.Errorf("Case %q: Output does not match expected:\n%s", testcase.name, diff.Diff(actualStdOut, expected))
+			}
+
+			if actualStdErr != "" {
+				t.Errorf("Case %q: Got error output:\n%s", testcase.name, actualStdErr)
+			}
+		})
+	}
+}
diff --git a/cli/cmds.go b/cli/cmds.go
index 6f09884b..fb4d8b5f 100644
--- a/cli/cmds.go
+++ b/cli/cmds.go
@@ -179,7 +179,7 @@ func Make() *cobra.Command {
 	diffCmd.Flags().StringP("pattern", "p", "", "glob pattern to match with each file name (e.g. *.markdown)")
 
 	// options
-	root.AddCommand(&cobra.Command{
+	blocksCmd := &cobra.Command{
 		Use:   "blocks [file]",
 		Short: "extracts terraform blocks from a file ",
 		//options: no header (######), format (json? xml? etc), only should block x?
@@ -193,10 +193,14 @@ func Make() *cobra.Command {
 			}
 			log.Debugf("terrafmt blocks %s", filename)
 
+			verbose := viper.GetBool("verbose")
+			zeroTerminated, _ := cmd.Flags().GetBool("zero-terminated")
 			fs := afero.NewOsFs()
-			return findBlocksInFile(fs, log, filename, viper.GetBool("verbose"), cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr())
+			return findBlocksInFile(fs, log, filename, verbose, zeroTerminated, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr())
 		},
-	})
+	}
+	root.AddCommand(blocksCmd)
+	blocksCmd.Flags().BoolP("zero-terminated", "0", false, "outputs blocks separated by null separator")
 
 	root.AddCommand(&cobra.Command{
 		Use:   "version",
@@ -303,15 +307,20 @@ func versionCmd(cmd *cobra.Command, args []string) {
 	fmt.Println("  + " + terraformVersion)
 }
 
-func findBlocksInFile(fs afero.Fs, log *logrus.Logger, filename string, verbose bool, stdin io.Reader, stdout, stderr io.Writer) error {
+func findBlocksInFile(fs afero.Fs, log *logrus.Logger, filename string, verbose, zeroTerminated bool, stdin io.Reader, stdout, stderr io.Writer) error {
 	br := blocks.Reader{
 		Log:      log,
 		ReadOnly: true,
 		LineRead: blocks.ReaderIgnore,
 		BlockRead: func(br *blocks.Reader, i int, b string) error {
 			outW := stdout
-			fmt.Fprint(outW, c.Sprintf("\n<white>#######</> <cyan>B%d</><darkGray> @ #%d</>\n", br.BlockCount, br.LineCount))
+			if !zeroTerminated {
+				fmt.Fprint(outW, c.Sprintf("\n<white>#######</> <cyan>B%d</><darkGray> @ #%d</>\n", br.BlockCount, br.LineCount))
+			}
 			fmt.Fprint(outW, b)
+			if zeroTerminated {
+				fmt.Fprint(outW, "\x00")
+			}
 			return nil
 		},
 	}