package main import ( "bufio" "context" "fmt" "os" "os/signal" "path/filepath" "regexp" "runtime" "strings" "sync" "sync/atomic" "syscall" "github.com/mattn/go-isatty" "github.com/schollz/progressbar/v3" "github.com/spf13/cobra" "github.com/ungluedlabs/go-jdenticon/jdenticon" ) const ( // Maximum filename length to avoid filesystem issues maxFilenameLength = 200 ) var ( // Keep a-z, A-Z, 0-9, underscore, hyphen, and period. Replace everything else. sanitizeRegex = regexp.MustCompile(`[^a-zA-Z0-9_.-]+`) ) // batchJob represents a single identicon generation job type batchJob struct { value string outputPath string size int } // batchStats tracks processing statistics atomically type batchStats struct { processed int64 failed int64 } // batchCmd represents the batch command var batchCmd = &cobra.Command{ Use: "batch ", Short: "Generate multiple identicons from a list using concurrent processing", Long: `Generate multiple identicons from a list of values in a text file. The input file should contain one value per line. Each value will be used to generate an identicon saved to the output directory. The filename will be based on the input value with special characters replaced. Uses concurrent processing with a worker pool for optimal performance. The number of concurrent workers can be controlled with the --concurrency flag. Examples: jdenticon batch users.txt --output-dir ./avatars jdenticon batch emails.txt --output-dir ./avatars --format png --size 64 jdenticon batch large-list.txt --output-dir ./avatars --concurrency 8`, Args: cobra.ExactArgs(1), // Validate argument count RunE: func(cmd *cobra.Command, args []string) error { return runConcurrentBatch(cmd, args) }, } func init() { rootCmd.AddCommand(batchCmd) // Define local flags specific to 'batch' batchCmd.Flags().StringP("output-dir", "d", "", "Output directory for generated identicons (required)") _ = batchCmd.MarkFlagRequired("output-dir") // Concurrency control batchCmd.Flags().IntP("concurrency", "c", runtime.NumCPU(), fmt.Sprintf("Number of concurrent workers (default: %d)", runtime.NumCPU())) } // runConcurrentBatch executes the batch processing with concurrent workers func runConcurrentBatch(cmd *cobra.Command, args []string) error { inputFile := args[0] outputDir, _ := cmd.Flags().GetString("output-dir") concurrency, _ := cmd.Flags().GetInt("concurrency") // Validate concurrency value if concurrency <= 0 { return fmt.Errorf("concurrency must be positive, got %d", concurrency) } // 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 } // Create generator with custom config and larger cache for concurrent workload cacheSize := generatorCacheSize * concurrency // Scale cache with worker count generator, err := jdenticon.NewGeneratorWithConfig(config, cacheSize) if err != nil { return fmt.Errorf("failed to create generator: %w", err) } // Create output directory if it doesn't exist // #nosec G301 -- 0755 is standard for directories; CLI tool needs world-readable output if err = os.MkdirAll(outputDir, 0755); err != nil { return fmt.Errorf("failed to create output directory: %w", err) } // Set up graceful shutdown context ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() // Process the file and collect jobs jobs, total, err := prepareJobs(inputFile, outputDir, format, size) if err != nil { return err } if total == 0 { fmt.Fprintf(os.Stderr, "No valid entries found in input file\n") return nil } // Initialize progress reporting showProgress := isTTY(os.Stderr) var bar *progressbar.ProgressBar if showProgress { bar = createProgressBar(total) } // Execute concurrent processing stats, err := processConcurrentJobs(ctx, jobs, generator, format, concurrency, bar) if err != nil { return err } // Final status message processed := atomic.LoadInt64(&stats.processed) failed := atomic.LoadInt64(&stats.failed) if showProgress { fmt.Fprintf(os.Stderr, "\nBatch processing complete: %d succeeded, %d failed\n", processed, failed) } else { fmt.Fprintf(os.Stderr, "Batch processing complete: %d succeeded, %d failed\n", processed, failed) } if failed > 0 { return fmt.Errorf("some identicons failed to generate") } return nil } // prepareJobs reads the input file and creates a slice of jobs func prepareJobs(inputFile, outputDir string, format FormatFlag, size int) ([]batchJob, int, error) { // #nosec G304 -- inputFile is provided via CLI flag, validated by cobra file, err := os.Open(inputFile) if err != nil { return nil, 0, fmt.Errorf("failed to open input file: %w", err) } defer file.Close() var jobs []batchJob scanner := bufio.NewScanner(file) for scanner.Scan() { value := strings.TrimSpace(scanner.Text()) if value == "" { continue // Skip empty lines } // Generate filename from value (sanitize for filesystem) filename := sanitizeFilename(value) extension := ".svg" if format == FormatPNG { extension = ".png" } outputPath := filepath.Join(outputDir, filename+extension) jobs = append(jobs, batchJob{ value: value, outputPath: outputPath, size: size, }) } if err := scanner.Err(); err != nil { return nil, 0, fmt.Errorf("error reading input file: %w", err) } return jobs, len(jobs), nil } // createProgressBar creates a progress bar for the batch processing func createProgressBar(total int) *progressbar.ProgressBar { return progressbar.NewOptions(total, progressbar.OptionSetWriter(os.Stderr), progressbar.OptionSetDescription("Processing identicons..."), progressbar.OptionEnableColorCodes(true), progressbar.OptionShowCount(), progressbar.OptionSetWidth(15), progressbar.OptionSetTheme(progressbar.Theme{ Saucer: "[green]=[reset]", SaucerHead: "[green]>[reset]", SaucerPadding: " ", BarStart: "[", BarEnd: "]", }), ) } // processConcurrentJobs executes jobs using a worker pool pattern func processConcurrentJobs(ctx context.Context, jobs []batchJob, generator *jdenticon.Generator, format FormatFlag, concurrency int, bar *progressbar.ProgressBar) (*batchStats, error) { stats := &batchStats{} jobChan := make(chan batchJob, len(jobs)) var wg sync.WaitGroup // Start worker goroutines for i := 0; i < concurrency; i++ { wg.Add(1) go func(workerID int) { defer wg.Done() batchWorker(ctx, workerID, jobChan, generator, format, stats, bar) }(i) } // Send jobs to workers go func() { defer close(jobChan) for _, job := range jobs { select { case jobChan <- job: case <-ctx.Done(): return } } }() // Wait for all workers to complete wg.Wait() // Check if processing was interrupted if ctx.Err() != nil { return stats, fmt.Errorf("processing interrupted: %w", ctx.Err()) } return stats, nil } // batchWorker processes jobs from the job channel func batchWorker(ctx context.Context, workerID int, jobs <-chan batchJob, generator *jdenticon.Generator, format FormatFlag, stats *batchStats, bar *progressbar.ProgressBar) { for { select { case job, ok := <-jobs: if !ok { // Channel closed, no more jobs return } // Process the job with context for cancellation support if err := processJob(ctx, job, generator, format); err != nil { fmt.Fprintf(os.Stderr, "Worker %d failed to process %q: %v\n", workerID, job.value, err) atomic.AddInt64(&stats.failed, 1) } else { atomic.AddInt64(&stats.processed, 1) } // Update progress bar if available if bar != nil { _ = bar.Add(1) } case <-ctx.Done(): // Shutdown signal received return } } } // processJob handles a single identicon generation job func processJob(ctx context.Context, job batchJob, generator *jdenticon.Generator, format FormatFlag) error { // Generate identicon with context for cancellation support icon, err := generator.Generate(ctx, job.value, job.size) if err != nil { return fmt.Errorf("failed to generate identicon: %w", err) } // Generate output based on format result, err := renderIcon(icon, format) if err != nil { return fmt.Errorf("failed to render %s: %w", format, err) } // Write file // #nosec G306 -- 0644 is appropriate for generated image files (world-readable) if err := os.WriteFile(job.outputPath, result, 0644); err != nil { return fmt.Errorf("failed to write file: %w", err) } return nil } // sanitizeFilename converts a value to a safe filename using a whitelist approach func sanitizeFilename(value string) string { // Replace @ separately for readability if desired filename := strings.ReplaceAll(value, "@", "_at_") // Replace all other invalid characters with a single underscore filename = sanitizeRegex.ReplaceAllString(filename, "_") // Limit length to avoid filesystem issues if len(filename) > maxFilenameLength { filename = filename[:maxFilenameLength] } return filename } // isTTY checks if the given file descriptor is a terminal func isTTY(f *os.File) bool { return isatty.IsTerminal(f.Fd()) }