Support plain user data

Author: reaandrew

Config.go

@@ -0,0 +1,29 @@
+package main
+
+import (
+	"fmt"
+)
+
+var providers = map[string]DataProvider{
+	"aws": NewDefaultAwsProvider(),
+}
+
+type Config struct {
+	outputDir   string
+	providerKey string
+	args        []string
+}
+
+func (config Config) getProvider() (DataProvider, error) {
+	if config.providerKey == "" {
+		input := config.args[0]
+		return CommandLineProvider{Input: input}, nil
+	} else {
+		var ok bool
+		provider, ok := providers[config.providerKey]
+		if !ok {
+			return nil, fmt.Errorf("unknown provider: %s", config.providerKey)
+		}
+		return provider, nil
+	}
+}

Config_test.go

@@ -0,0 +1,11 @@
+package main
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestHandlesUnknownProvider(t *testing.T) {
+	_, err := Config{providerKey: "talula"}.getProvider()
+	assert.ErrorContains(t, err, "unknown provider")
+}

InputProcessor.go

@@ -0,0 +1,51 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+)
+
+type InputProcessor struct {
+	config Config
+}
+
+func (inputProcessor InputProcessor) writePlainUserDataFile(input DataOutputPair) error {
+	outputPath := filepath.Join(inputProcessor.config.outputDir, input.OutputDir)
+	fullPath := filepath.Join(outputPath, "userdata")
+	err := os.MkdirAll(path.Dir(fullPath), 0755)
+
+	if err != nil {
+		return fmt.Errorf("error creating output directories: %w", err)
+	}
+	err = os.WriteFile(fullPath, input.Data, 0755)
+	if err != nil {
+		return fmt.Errorf("errir writing file: %w", err)
+	}
+	return nil
+}
+
+func (inputProcessor InputProcessor) Process(inputs []DataOutputPair) error {
+	for _, input := range inputs {
+		decoded, err := decode(input.Data)
+
+		if err != nil {
+			return inputProcessor.writePlainUserDataFile(input)
+		} else {
+			attachments, err := ExtractMimeAttachmentsFromBytes(decoded)
+			if err != nil {
+				return inputProcessor.writePlainUserDataFile(input)
+			} else {
+				outputPath := filepath.Join(inputProcessor.config.outputDir, input.OutputDir)
+
+				err = ExtractCloudConfig(attachments, outputPath)
+				if err != nil {
+					return fmt.Errorf("failed to save write files: %w", err)
+				}
+			}
+		}
+	}
+
+	return nil
+}

InputProcessor_test.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWritePlainUserDataFileWhenNotBase64Encoded(t *testing.T) {
+	tempDir, _ := os.MkdirTemp("", "test")
+	defer os.RemoveAll(tempDir)
+
+	config := Config{outputDir: tempDir}
+	inputProcessor := InputProcessor{config: config}
+
+	dataOutputPair := DataOutputPair{
+		Data:      []byte("some data"),
+		OutputDir: "some_output_dir",
+	}
+
+	err := inputProcessor.writePlainUserDataFile(dataOutputPair)
+
+	assert.Nil(t, err)
+
+	expectedFilePath := filepath.Join(tempDir, "some_output_dir", "userdata")
+	_, err = os.Stat(expectedFilePath)
+	assert.False(t, os.IsNotExist(err))
+}
+
+func TestWritePlainUserDataFileWhenNotMime(t *testing.T) {
+	tempDir, _ := os.MkdirTemp("", "test")
+	defer os.RemoveAll(tempDir)
+
+	config := Config{outputDir: tempDir}
+	inputProcessor := InputProcessor{config: config}
+
+	dataOutputPairs := []DataOutputPair{
+		{
+			Data:      []byte(base64EncodedHelloWorld),
+			OutputDir: "some_output_dir1",
+		},
+	}
+
+	err := inputProcessor.Process(dataOutputPairs)
+
+	assert.Nil(t, err)
+}
+
+func TestProcess(t *testing.T) {
+
+	tempDir, _ := os.MkdirTemp("", "test")
+	defer os.RemoveAll(tempDir)
+
+	config := Config{outputDir: tempDir}
+	inputProcessor := InputProcessor{config: config}
+
+	dataOutputPairs := []DataOutputPair{
+		{
+			Data:      []byte("some data"),
+			OutputDir: "some_output_dir1",
+		},
+		{
+			Data:      []byte("some other data"),
+			OutputDir: "some_output_dir2",
+		},
+	}
+
+	err := inputProcessor.Process(dataOutputPairs)
+
+	assert.Nil(t, err)
+}

Mime.go

@@ -4,13 +4,16 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/base64"
+	"errors"
 	"io"
 	"log"
 	"mime"
 	"mime/multipart"
 	"strings"
 )
 
+var errFailedToExtractMimeBoundary = errors.New("failed to get MIME boundary")
+
 type MimeAttachment struct {
 	ContentType string
 	Content     []byte
@@ -35,10 +38,10 @@ func extractBoundary(data []byte) (string, error) {
 	return "", io.EOF
 }
 
-func decodeMimAttachments(data []byte) (attachments []MimeAttachment, err error) {
+func ExtractMimeAttachmentsFromBytes(data []byte) (attachments []MimeAttachment, err error) {
 	boundary, err := extractBoundary(data)
 	if err != nil {
-		log.Fatalf("Failed to get boundary: %s", err)
+		return []MimeAttachment{}, errFailedToExtractMimeBoundary
 	}
 
 	reader := multipart.NewReader(bytes.NewReader(data), boundary)
@@ -81,13 +84,3 @@ func decodeMimAttachments(data []byte) (attachments []MimeAttachment, err error)
 
 	return attachments, err
 }
-
-func ExtractMimeAttachmentsFromBytes(encodedData []byte) (attachments []MimeAttachment, err error) {
-	decoded, err := decode(encodedData)
-	if err != nil {
-
-		return
-	}
-
-	return decodeMimAttachments(decoded)
-}

Mime_test.go

@@ -0,0 +1,24 @@
+package main
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestExtractMimeAttachments(t *testing.T) {
+	decoded, _ := decode([]byte(mimeMessage))
+
+	attachments, err := ExtractMimeAttachmentsFromBytes(decoded)
+
+	assert.Nil(t, err)
+	assert.Len(t, attachments, 2)
+	assert.Equal(t, attachments[0].ContentType, "text/cloud-config; charset=\"utf-8\"")
+	assert.Equal(t, attachments[1].ContentType, "text/x-shellscript; charset=\"utf-8\"")
+}
+
+func TestExtractMimeAttachmentsReturnsErrorWhenNotMime(t *testing.T) {
+	_, err := ExtractMimeAttachmentsFromBytes([]byte("not mime"))
+
+	assert.NotNil(t, err)
+	assert.Equal(t, err, errFailedToExtractMimeBoundary)
+}

main.go

@@ -4,35 +4,10 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"log"
 	"os"
-	"path/filepath"
 )
 
-var providers = map[string]DataProvider{
-	"aws": NewDefaultAwsProvider(),
-}
-
-type Config struct {
-	outputDir   string
-	providerKey string
-	args        []string
-}
-
-func (config Config) getProvider() (DataProvider, error) {
-	if config.providerKey == "" {
-		input := config.args[0]
-		return CommandLineProvider{Input: input}, nil
-	} else {
-		var ok bool
-		provider, ok := providers[config.providerKey]
-		if !ok {
-			fmt.Println("Unknown provider:", config.providerKey)
-			os.Exit(1)
-		}
-		return provider, nil
-	}
-}
-
 func main() {
 	if err := run(); err != nil {
 		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
@@ -48,29 +23,17 @@ func run() error {
 	}
 	provider, err := config.getProvider()
 	if err != nil {
-		return err
+		log.Fatalf("failed to the provider")
 	}
 
 	inputs, err := provider.FetchData()
 	if err != nil {
 		return fmt.Errorf("failed to fetch data: %w", err)
 	}
 
-	for _, input := range inputs {
-		attachments, err := ExtractMimeAttachmentsFromBytes(input.Data)
-		if err != nil {
-			return fmt.Errorf("failed to extract mime attachments: %w", err)
-		}
-
-		outputPath := filepath.Join(config.outputDir, input.OutputDir)
-
-		err = ExtractCloudConfig(attachments, outputPath)
-		if err != nil {
-			return fmt.Errorf("failed to save write files: %w", err)
-		}
-	}
-
-	return nil
+	return InputProcessor{
+		config: config,
+	}.Process(inputs)
 }
 
 func parseFlags() (config Config, err error) {

main_test.go

@@ -1,10 +1,5 @@
 package main
 
-import (
-	"github.com/stretchr/testify/assert"
-	"testing"
-)
-
 const (
 	base64EncodedHelloWorld       = "SGVsbG8sIFdvcmxkIQ=="
 	base64EncodedZippedHelloWorld = "H4sIAAAAAAAAA/NIzcnJ11EIzy/KSVEEANDDSuwNAAAA"
@@ -47,13 +42,3 @@ const (
 	mimeMessage = `
 H4sICFuIGWUAA3VzZXJkYXRhAL2USXOjOBTH73wKV+5OC7DTISkfzCKMbXAjQCw3IQhgi2UMXuDTtzqdSk1P9aGnpmp0kepJb/m/+ulpbTPkzTD3xy5/mdUXNlQdOQ9f6uqeZ6+ztL00GTmPq4fVr2u5kL8+K8ulAhaKoiyWXyWwWj0ItmUbc5yf+6ptXmbiIxCE+fxPXAXtl0KG/D58oay9ZHPaNm9V8TqjJTn3+bB6uAxv8+ffpfoMcSZN/5af50ZD26xqipdZSvr8afH5Qq/6ru2r4d2VDAOhZc3tr7O3iuUNqfPVw9+TP46kZg+CYI3qMTUVMdHEYyotaxLSSx7CIV13Mq0ZSLDCbXdGp3ZnaeujBbORRIhZG9Qmnvrhk4m0hkAgoXKxzKRLzaDIzOdiL2EQj8oUh+ItNfFbHKm3/a0tLA0U1ISA6PxcKSwz7SuV4EDNO4skeKOaMgmpJLJ9xMo0vP2sSLN+VFBYJr7EknIi4bI5VGqTc83O0f24c65pg1jauE+WUQAKmO6KCRTQKcNuYEx+k1RYxFxWZmbmVnd8nOAJTm4QjIEhoj1wtu7JMQ/es+QBW6YAASRvgXPKbrYhXpC/vgs2YHdPxDtkwirUlDiV1Q5JCrGl5NtuYvoh6o4pKBcJ65AbbWVvygJ0Go4xRgHZlJV7zM426HwCMRb2Ipbsk2MEG3i39ZNEwuEebnBzgEnPK/So3MVpfT9kJ2gjyR1DYwuJATUHWKMPtnqiJ3IsBpJvxkshCcoxj7o2xKqPpFL3ZTSkOBETM9kTyHyiOzd/Kv2MGcCZLJHqaPQDe4kC8eTrpYYmqPpNCWgAOwGHzmYHysrb2KIfLg+pHtyIr9bIHBYcmUsKsRls4ttuVPY7ADuygUdvKjtPLpaBBBcx6/pDAPcZs0fBa8pvLoaq6/6WI5C8IwIBP7/zxu1XS1NEKlscjbKkdfCJiPCTkecfgQZro5aZWTxZmnXNImfcy04bR1vG9zKN1J4DXNLNmjsrQ+KJIIlKsG/QIvvASfjXPEV4i8R/4GRkoWMiXXB8lXkQqt5UyAFEIQ4VP/bVK+/lXz4OFsRgNreZbuAwu062Id6qbgBVrejGLFwe09B90qp1sffUVOBf410/b/YxDm/XuGYXLgtweVyC01oMrP7LMLrP+zJnrKfnqhv+92nUD3woP/YlH0JJ6LTpuK4CGZWUtzvRrD8WNp8LwncQdAAY8gUAAA==`
 )
-
-func TestExtractMimeAttachments(t *testing.T) {
-
-	attachments, err := ExtractMimeAttachmentsFromBytes([]byte(mimeMessage))
-
-	assert.Nil(t, err)
-	assert.Len(t, attachments, 2)
-	assert.Equal(t, attachments[0].ContentType, "text/cloud-config; charset=\"utf-8\"")
-	assert.Equal(t, attachments[1].ContentType, "text/x-shellscript; charset=\"utf-8\"")
-}