router-backup/src/main.go

// Copyright 2025, IPng Networks GmbH, Pim van Pelt <pim@ipng.ch>

package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"

	"github.com/spf13/cobra"
)

const Version = "1.3.0"

func processDevice(hostname string, deviceConfig Device, commands []string, excludePatterns []string, password, keyFile string, port int, outputDir string) bool {
	// Create backup instance
	backup := NewRouterBackup(hostname, deviceConfig.Address, deviceConfig.User, password, keyFile, port)

	// Connect and backup
	if err := backup.Connect(); err != nil {
		fmt.Printf("%s: Failed to connect: %v\n", hostname, err)
		return false
	}

	err := backup.BackupCommands(commands, excludePatterns, outputDir)
	backup.Disconnect()

	if err != nil {
		fmt.Printf("%s: Backup failed: %v\n", hostname, err)
		return false
	} else {
		fmt.Printf("%s: Backup completed\n", hostname)
		return true
	}
}

func main() {
	var yamlFiles []string
	var password string
	var keyFile string
	var port int
	var outputDir string
	var hostFilter []string
	var parallel int

	var rootCmd = &cobra.Command{
		Use:     "ipng-router-backup",
		Short:   "SSH Router Backup Tool",
		Long:    "Connects to routers via SSH and runs commands, saving output to local files.",
		Version: Version,
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("IPng Networks Router Backup v%s\n", Version)

			// Expand glob patterns in YAML files
			var expandedYamlFiles []string
			for _, pattern := range yamlFiles {
				matches, err := filepath.Glob(pattern)
				if err != nil {
					log.Fatalf("Invalid glob pattern '%s': %v", pattern, err)
				}
				if len(matches) == 0 {
					log.Fatalf("No files matched pattern '%s'", pattern)
				}
				expandedYamlFiles = append(expandedYamlFiles, matches...)
			}

			// Load configuration
			cfg, err := ConfigRead(expandedYamlFiles)
			if err != nil {
				log.Fatalf("Failed to load config: %v", err)
			}

			// Check authentication setup
			hasAuth := 0
			if os.Getenv("SSH_AUTH_SOCK") != "" {
				fmt.Println("Using SSH agent for authentication")
				hasAuth++
			}
			if keyFile == "" {
				keyFile = findDefaultSSHKey()
				if keyFile != "" {
					fmt.Printf("Using SSH key: %s\n", keyFile)
					hasAuth++
				}
			}
			if password != "" {
				fmt.Println("Using --password for authentication")
				hasAuth++
			}
			if hasAuth == 0 {
				log.Fatal("No authentication mechanisms found.")
			}

			// Process devices
			if len(cfg.Devices) == 0 {
				log.Fatal("No devices found in config file")
			}

			// Filter devices if --host flags are provided
			devicesToProcess := cfg.Devices
			if len(hostFilter) > 0 {
				devicesToProcess = make(map[string]Device)
				for _, pattern := range hostFilter {
					patternMatched := false
					for hostname, deviceConfig := range cfg.Devices {
						if matched, _ := filepath.Match(pattern, hostname); matched {
							devicesToProcess[hostname] = deviceConfig
							patternMatched = true
						}
					}
					if !patternMatched {
						fmt.Printf("Warning: Host pattern '%s' did not match any devices\n", pattern)
					}
				}
			}

			totalCount := len(devicesToProcess)

			// Create channels for work distribution and result collection
			type DeviceWork struct {
				hostname        string
				deviceConfig    Device
				commands        []string
				excludePatterns []string
			}

			type DeviceResult struct {
				hostname string
				success  bool
			}

			workChan := make(chan DeviceWork, totalCount)
			resultChan := make(chan DeviceResult, totalCount)

			// Start worker pool
			var wg sync.WaitGroup
			for i := 0; i < parallel; i++ {
				wg.Add(1)
				go func() {
					defer wg.Done()
					for work := range workChan {
						fmt.Printf("%s: Processing device (type: %s)\n", work.hostname, work.deviceConfig.Type)

						success := processDevice(work.hostname, work.deviceConfig, work.commands, work.excludePatterns, password, keyFile, port, outputDir)
						resultChan <- DeviceResult{hostname: work.hostname, success: success}
					}
				}()
			}

			// Queue all work
			for hostname, deviceConfig := range devicesToProcess {
				user := deviceConfig.User
				commands := deviceConfig.Commands
				deviceType := deviceConfig.Type
				var excludePatterns []string

				// If device has a type, get commands and exclude patterns from types section
				if deviceType != "" {
					if typeConfig, exists := cfg.Types[deviceType]; exists {
						commands = typeConfig.Commands
						excludePatterns = typeConfig.Exclude
					}
				}

				if user == "" {
					fmt.Printf("%s: No user specified, skipping\n", hostname)
					continue
				}

				if len(commands) == 0 {
					fmt.Printf("%s: No commands specified, skipping\n", hostname)
					continue
				}

				workChan <- DeviceWork{
					hostname:        hostname,
					deviceConfig:    deviceConfig,
					commands:        commands,
					excludePatterns: excludePatterns,
				}
			}
			close(workChan)

			// Wait for all workers to finish
			go func() {
				wg.Wait()
				close(resultChan)
			}()

			// Collect results
			successCount := 0
			for result := range resultChan {
				if result.success {
					successCount++
				}
			}

			fmt.Printf("Overall summary: %d/%d devices processed successfully\n", successCount, totalCount)

			// Set exit code based on results
			if successCount == 0 {
				os.Exit(11) // All devices failed
			} else if successCount < totalCount {
				os.Exit(10) // Some devices failed
			}
			// Exit code 0 (success) when all devices succeeded
		},
	}

	rootCmd.Flags().StringSliceVar(&yamlFiles, "yaml", []string{}, "YAML configuration file paths (required, can be repeated)")
	rootCmd.Flags().StringVar(&password, "password", "", "SSH password")
	rootCmd.Flags().StringVar(&keyFile, "key-file", "", "SSH private key file path")
	rootCmd.Flags().IntVar(&port, "port", 22, "SSH port")
	rootCmd.Flags().StringVar(&outputDir, "output-dir", "/tmp", "Output directory for command output files")
	rootCmd.Flags().StringSliceVar(&hostFilter, "host", []string{}, "Specific host(s) to process (can be repeated, processes all if not specified)")
	rootCmd.Flags().IntVar(&parallel, "parallel", 10, "Maximum number of devices to process in parallel")

	if err := rootCmd.MarkFlagRequired("yaml"); err != nil {
		log.Fatal(err)
	}

	if err := rootCmd.Execute(); err != nil {
		log.Fatal(err)
	}
}