go-jdenticon/cmd/jdenticon/generate_test.go

package main

import (
	"bytes"
	"context"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/ungluedlabs/go-jdenticon/jdenticon"
)

// TestGenerateCommand tests the generate command functionality
func TestGenerateCommand(t *testing.T) {
	tests := []struct {
		name        string
		args        []string
		wantErr     bool
		outputCheck func(t *testing.T, output []byte, outputFile string)
	}{
		{
			name:    "generate SVG to stdout",
			args:    []string{"generate", "test@example.com"},
			wantErr: false,
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				if !bytes.Contains(output, []byte("<svg")) {
					t.Error("Expected SVG output to contain <svg tag")
				}
				if !bytes.Contains(output, []byte("xmlns=\"http://www.w3.org/2000/svg\"")) {
					t.Error("Expected SVG output to contain namespace declaration")
				}
			},
		},
		{
			name:    "generate with custom size",
			args:    []string{"generate", "--size", "128", "test@example.com"},
			wantErr: false,
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				if !bytes.Contains(output, []byte("width=\"128\"")) {
					t.Error("Expected SVG to have width=128")
				}
				if !bytes.Contains(output, []byte("height=\"128\"")) {
					t.Error("Expected SVG to have height=128")
				}
			},
		},
		{
			name:    "generate PNG format",
			args:    []string{"generate", "--format", "png", "test@example.com"},
			wantErr: false,
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				// PNG files start with specific magic bytes
				if len(output) < 8 || !bytes.Equal(output[:8], []byte{137, 80, 78, 71, 13, 10, 26, 10}) {
					t.Error("Expected PNG output to have PNG magic bytes")
				}
			},
		},
		{
			name:    "generate with background color",
			args:    []string{"generate", "--bg-color", "#ffffff", "test@example.com"},
			wantErr: false,
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				// SVG should contain background rect
				if !bytes.Contains(output, []byte("fill=\"#ffffff\"")) {
					t.Error("Expected SVG to contain background color")
				}
			},
		},
		{
			name:    "generate with custom padding",
			args:    []string{"generate", "--padding", "0.2", "test@example.com"},
			wantErr: false,
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				// Should generate valid SVG
				if !bytes.Contains(output, []byte("<svg")) {
					t.Error("Expected valid SVG output")
				}
			},
		},
		{
			name:    "no arguments",
			args:    []string{},
			wantErr: false, // Shows help, doesn't error
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				// Should show help text
				outputStr := string(output)
				if !strings.Contains(outputStr, "Usage:") && !strings.Contains(outputStr, "generate") {
					t.Error("Expected help text to be shown when no arguments provided")
				}
			},
		},
		{
			name:    "too many arguments",
			args:    []string{"arg1", "arg2"},
			wantErr: true,
			outputCheck: func(t *testing.T, output []byte, outputFile string) {
				// Should not produce output on error
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Reset viper for clean state
			viper.Reset()

			// Create a buffer to capture output
			var output bytes.Buffer

			// Create the generate command for testing
			cmd := createTestGenerateCommand()
			cmd.SetOut(&output)
			cmd.SetErr(&output)

			// Set args and execute
			cmd.SetArgs(tt.args)
			err := cmd.Execute()

			if (err != nil) != tt.wantErr {
				t.Errorf("generateCmd.Execute() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			if !tt.wantErr && tt.outputCheck != nil {
				tt.outputCheck(t, output.Bytes(), "")
			}
		})
	}
}

// TestGenerateToFile tests file output functionality
func TestGenerateToFile(t *testing.T) {
	// Create a temporary directory for test outputs
	tempDir, err := os.MkdirTemp("", "jdenticon-test")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tempDir)

	// Change to temp directory to test file creation there
	originalDir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Failed to get current directory: %v", err)
	}
	defer os.Chdir(originalDir)

	if err := os.Chdir(tempDir); err != nil {
		t.Fatalf("Failed to change to temp directory: %v", err)
	}

	tests := []struct {
		name      string
		args      []string
		filename  string
		wantErr   bool
		fileCheck func(t *testing.T, filepath string)
	}{
		{
			name:     "generate SVG to file",
			args:     []string{"generate", "test@example.com"},
			filename: "test.svg",
			wantErr:  false,
			fileCheck: func(t *testing.T, filepath string) {
				content, err := os.ReadFile(filepath)
				if err != nil {
					t.Fatalf("Failed to read output file: %v", err)
				}
				if !bytes.Contains(content, []byte("<svg")) {
					t.Error("Expected SVG file to contain <svg tag")
				}
			},
		},
		{
			name:     "generate PNG to file",
			args:     []string{"generate", "--format", "png", "test@example.com"},
			filename: "test.png",
			wantErr:  false,
			fileCheck: func(t *testing.T, filepath string) {
				content, err := os.ReadFile(filepath)
				if err != nil {
					t.Fatalf("Failed to read output file: %v", err)
				}
				// Check PNG magic bytes
				if len(content) < 8 || !bytes.Equal(content[:8], []byte{137, 80, 78, 71, 13, 10, 26, 10}) {
					t.Error("Expected PNG file to have PNG magic bytes")
				}
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			viper.Reset()

			// Use relative path since we're in temp directory
			outputPath := tt.filename
			args := append(tt.args, "--output", outputPath)

			var output bytes.Buffer
			cmd := createTestGenerateCommand()
			cmd.SetOut(&output)
			cmd.SetErr(&output)
			cmd.SetArgs(args)

			err := cmd.Execute()

			if (err != nil) != tt.wantErr {
				t.Errorf("generateCmd.Execute() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			if !tt.wantErr {
				// Check that file was created
				if _, err := os.Stat(outputPath); os.IsNotExist(err) {
					t.Errorf("Expected output file to be created at %s", outputPath)
					return
				}

				if tt.fileCheck != nil {
					tt.fileCheck(t, outputPath)
				}
			}
		})
	}
}

// TestGenerateValidation tests input validation
func TestGenerateValidation(t *testing.T) {
	tests := []struct {
		name       string
		args       []string
		wantErr    bool
		errorCheck func(t *testing.T, err error)
	}{
		{
			name:    "negative size",
			args:    []string{"generate", "--size", "-1", "test@example.com"},
			wantErr: true,
			errorCheck: func(t *testing.T, err error) {
				if !strings.Contains(err.Error(), "size must be positive") {
					t.Errorf("Expected size validation error, got: %v", err)
				}
			},
		},
		{
			name:    "zero size",
			args:    []string{"generate", "--size", "0", "test@example.com"},
			wantErr: true,
			errorCheck: func(t *testing.T, err error) {
				if !strings.Contains(err.Error(), "size must be positive") {
					t.Errorf("Expected size validation error, got: %v", err)
				}
			},
		},
		{
			name:    "invalid padding",
			args:    []string{"generate", "--padding", "-0.1", "test@example.com"},
			wantErr: true,
			errorCheck: func(t *testing.T, err error) {
				if !strings.Contains(err.Error(), "padding") {
					t.Errorf("Expected padding validation error, got: %v", err)
				}
			},
		},
		{
			name:    "invalid color saturation",
			args:    []string{"generate", "--color-saturation", "1.5", "test@example.com"},
			wantErr: true,
			errorCheck: func(t *testing.T, err error) {
				if !strings.Contains(err.Error(), "saturation") {
					t.Errorf("Expected saturation validation error, got: %v", err)
				}
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			viper.Reset()

			var output bytes.Buffer
			cmd := createTestGenerateCommand()
			cmd.SetOut(&output)
			cmd.SetErr(&output)
			cmd.SetArgs(tt.args)

			err := cmd.Execute()

			if (err != nil) != tt.wantErr {
				t.Errorf("generateCmd.Execute() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			if tt.wantErr && tt.errorCheck != nil {
				tt.errorCheck(t, err)
			}
		})
	}
}

// TestPathTraversalSecurity tests the path traversal vulnerability fix
func TestPathTraversalSecurity(t *testing.T) {
	// Create a temporary directory for test outputs
	tempDir, err := os.MkdirTemp("", "jdenticon-security-test")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tempDir)

	// Change to temp directory to have a controlled test environment
	originalDir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Failed to get current directory: %v", err)
	}
	defer os.Chdir(originalDir)

	if err := os.Chdir(tempDir); err != nil {
		t.Fatalf("Failed to change to temp directory: %v", err)
	}

	// Create a subdirectory to test valid subdirectory writes
	subDir := filepath.Join(tempDir, "images")
	if err := os.Mkdir(subDir, 0755); err != nil {
		t.Fatalf("Failed to create subdirectory: %v", err)
	}

	tests := []struct {
		name         string
		outputPath   string
		expectError  bool
		errorMessage string
		description  string
	}{
		{
			name:        "valid_current_dir",
			outputPath:  "avatar.png",
			expectError: false,
			description: "Should allow files in current directory",
		},
		{
			name:        "valid_subdirectory",
			outputPath:  "images/avatar.png",
			expectError: false,
			description: "Should allow files in subdirectories",
		},
		{
			name:        "valid_relative_current",
			outputPath:  "./avatar.png",
			expectError: false,
			description: "Should allow explicit current directory notation",
		},
		{
			name:         "path_traversal_up_one",
			outputPath:   "../avatar.png",
			expectError:  true,
			errorMessage: "outside allowed directory",
			description:  "Should block path traversal attempts with ../",
		},
		{
			name:         "path_traversal_up_multiple",
			outputPath:   "../../avatar.png",
			expectError:  true,
			errorMessage: "outside allowed directory",
			description:  "Should block multiple directory traversal attempts",
		},
		{
			name:         "path_traversal_complex",
			outputPath:   "../../../etc/passwd",
			expectError:  true,
			errorMessage: "outside allowed directory",
			description:  "Should block attempts to write to system directories",
		},
		{
			name:         "absolute_path_system",
			outputPath:   filepath.Join(os.TempDir(), "avatar.png"),
			expectError:  true,
			errorMessage: "outside allowed directory",
			description:  "Should block absolute paths to system directories",
		},
		{
			name: "absolute_path_root",
			outputPath: func() string {
				if runtime.GOOS == "windows" {
					return `C:\Windows\test.png`
				} else {
					return "/etc/passwd"
				}
			}(),
			expectError:  true,
			errorMessage: "outside allowed directory",
			description:  "Should block attempts to overwrite system files",
		},
		{
			name:         "mixed_path_traversal",
			outputPath:   filepath.Join(".", "images", "..", "..", "..", os.TempDir(), "avatar.png"),
			expectError:  true,
			errorMessage: "outside allowed directory",
			description:  "Should block mixed path traversal attempts",
		},
		{
			name:         "windows_style_traversal",
			outputPath:   "..\\..\\evil.exe",
			expectError:  runtime.GOOS == "windows", // Only expect error on Windows
			errorMessage: "outside allowed directory",
			description:  "Should block Windows-style path traversal on Windows (backslashes are valid filename chars on Unix)",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			viper.Reset()

			// Test args for generating an identicon with the specified output path
			args := []string{"generate", "--output", tt.outputPath, "test@example.com"}

			var output bytes.Buffer
			cmd := createTestGenerateCommand()
			cmd.SetOut(&output)
			cmd.SetErr(&output)
			cmd.SetArgs(args)

			err := cmd.Execute()

			if tt.expectError {
				if err == nil {
					t.Errorf("Test %s: Expected error but got none. %s", tt.name, tt.description)
					return
				}
				if tt.errorMessage != "" && !strings.Contains(err.Error(), tt.errorMessage) {
					t.Errorf("Test %s: Expected error to contain %q, got: %v", tt.name, tt.errorMessage, err)
				}
				t.Logf("Test %s: Correctly blocked path %q with error: %v", tt.name, tt.outputPath, err)
			} else {
				if err != nil {
					t.Errorf("Test %s: Expected no error but got: %v. %s", tt.name, err, tt.description)
					return
				}

				// For successful cases, verify the file was created in the expected location
				expectedPath := filepath.Join(tempDir, tt.outputPath)
				expectedPath = filepath.Clean(expectedPath)

				if _, err := os.Stat(expectedPath); os.IsNotExist(err) {
					t.Errorf("Test %s: Expected file to be created at %s", tt.name, expectedPath)
				} else {
					t.Logf("Test %s: Successfully created file at %s", tt.name, expectedPath)
					// Clean up the created file
					os.Remove(expectedPath)
				}
			}
		})
	}
}

// TestValidateAndResolveOutputPath tests the security helper function directly
func TestValidateAndResolveOutputPath(t *testing.T) {
	tempDir, err := os.MkdirTemp("", "jdenticon-path-test")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tempDir)

	// Change to temp directory to test relative path resolution correctly
	originalDir, err := os.Getwd()
	if err != nil {
		t.Fatalf("Failed to get current directory: %v", err)
	}
	defer os.Chdir(originalDir)

	if err := os.Chdir(tempDir); err != nil {
		t.Fatalf("Failed to change to temp directory: %v", err)
	}

	tests := []struct {
		name        string
		baseDir     string
		outputPath  string
		expectError bool
		description string
	}{
		{
			name:        "valid_relative_file",
			baseDir:     tempDir,
			outputPath:  "test.png",
			expectError: false,
			description: "Valid relative file path should be allowed",
		},
		{
			name:        "valid_subdirectory_file",
			baseDir:     tempDir,
			outputPath:  "sub/test.png",
			expectError: false,
			description: "Valid file in subdirectory should be allowed",
		},
		{
			name:        "traversal_attack",
			baseDir:     tempDir,
			outputPath:  "../../../etc/passwd",
			expectError: true,
			description: "Path traversal attack should be blocked",
		},
		{
			name:        "absolute_outside_path",
			baseDir:     tempDir,
			outputPath:  filepath.Join(os.TempDir(), "test.png"),
			expectError: true,
			description: "Absolute path outside base should be blocked",
		},
		{
			name:        "current_dir_notation",
			baseDir:     tempDir,
			outputPath:  "./test.png",
			expectError: false,
			description: "Current directory notation should be allowed",
		},
		{
			name:        "complex_traversal",
			baseDir:     tempDir,
			outputPath:  "sub/../../../secret.txt",
			expectError: true,
			description: "Complex path traversal should be blocked",
		},
		{
			name:        "absolute_inside_allowed",
			baseDir:     tempDir,
			outputPath:  filepath.Join(tempDir, "allowed.png"),
			expectError: false,
			description: "Absolute path inside base directory should be allowed",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result, err := validateAndResolveOutputPath(tt.baseDir, tt.outputPath)

			if tt.expectError {
				if err == nil {
					t.Errorf("Expected error for %s, but got none. Result: %s", tt.description, result)
				}
			} else {
				if err != nil {
					t.Errorf("Unexpected error for %s: %v", tt.description, err)
				} else {
					// Verify the result is within the base directory
					// Use EvalSymlinks to handle macOS symlink issues
					absBase, _ := filepath.Abs(tt.baseDir)
					resolvedBase, err := filepath.EvalSymlinks(absBase)
					if err != nil {
						resolvedBase = absBase // fallback to original if symlink resolution fails
					}

					resolvedResult, err := filepath.EvalSymlinks(filepath.Dir(result))
					if err != nil {
						resolvedResult = filepath.Dir(result) // fallback to original if symlink resolution fails
					}
					resolvedResult = filepath.Join(resolvedResult, filepath.Base(result))

					if !strings.HasPrefix(resolvedResult, resolvedBase) {
						t.Errorf("Result path %s (resolved: %s) is not within base directory %s (resolved: %s)", result, resolvedResult, absBase, resolvedBase)
					}
				}
			}
		})
	}
}

// createTestGenerateCommand creates a generate command for testing
func createTestGenerateCommand() *cobra.Command {
	// Create root command with flags
	rootCmd := &cobra.Command{
		Use:   "jdenticon",
		Short: "Generate identicons from any input string",
	}

	// Initialize root flags
	initTestFlags(rootCmd)

	// Create generate command similar to the actual one
	generateCmd := &cobra.Command{
		Use:   "generate <value>",
		Short: "Generate a single identicon",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			// Get the input value
			value := args[0]

			// Get output file flag
			outputFile, _ := cmd.Flags().GetString("output")

			// Get format from viper
			format, err := getFormatFromViper()
			if err != nil {
				return err
			}

			// Populate library config from root persistent flags
			config, size, err := populateConfigFromFlags()
			if err != nil {
				return err
			}

			// Generate the identicon with custom config
			generator, err := jdenticon.NewGeneratorWithConfig(config, generatorCacheSize)
			if err != nil {
				return err
			}

			icon, err := generator.Generate(context.Background(), value, size)
			if err != nil {
				return err
			}

			// Generate output based on format
			result, err := renderIcon(icon, format)
			if err != nil {
				return err
			}

			// Output to file or stdout
			if outputFile != "" {
				// Determine the base directory for allowed writes. For a CLI, this is typically the CWD.
				baseDir, err := os.Getwd()
				if err != nil {
					return err
				}

				// Validate and resolve the user-provided output path.
				safeOutputPath, err := validateAndResolveOutputPath(baseDir, outputFile)
				if err != nil {
					// This is a security-related error, explicitly state it.
					return err
				}

				// Now use the safe and validated path for writing.
				if err := os.WriteFile(safeOutputPath, result, 0644); err != nil {
					return err
				}
			} else {
				// Write to stdout for piping
				if _, err := cmd.OutOrStdout().Write(result); err != nil {
					return err
				}
			}

			return nil
		},
	}

	// Add generate-specific flags
	generateCmd.Flags().StringP("output", "o", "", "Output file path. If empty, writes to stdout.")

	// Add to root command
	rootCmd.AddCommand(generateCmd)

	return rootCmd
}