From 64a4d53cf72dbdf725a958ac53d9aed6ec91996a Mon Sep 17 00:00:00 2001 From: Pim van Pelt Date: Sat, 5 Jul 2025 23:38:28 +0000 Subject: [PATCH] Rewrite Python to Go --- .gitignore | 1 + requirements.txt | 1 - router_backup.py | 235 -------------------------------- src/go.mod | 15 +++ src/go.sum | 19 +++ src/router_backup.go | 315 +++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 350 insertions(+), 236 deletions(-) create mode 100644 .gitignore delete mode 100644 requirements.txt delete mode 100755 router_backup.py create mode 100644 src/go.mod create mode 100644 src/go.sum create mode 100644 src/router_backup.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..6e0b016 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +router_backup diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 88fdc29..0000000 --- a/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -paramiko>=2.9.0 \ No newline at end of file diff --git a/router_backup.py b/router_backup.py deleted file mode 100755 index 1b68ff5..0000000 --- a/router_backup.py +++ /dev/null @@ -1,235 +0,0 @@ -#!/usr/bin/env python3 -""" -SSH Router Backup Script -Connects to routers via SSH and runs commands, saving output to local files. -""" - -import paramiko -import sys -import os -from datetime import datetime -import argparse -import json -import yaml - - -class RouterBackup: - def __init__(self, hostname, username, password=None, key_file=None, port=22): - self.hostname = hostname - self.username = username - self.password = password - self.key_file = key_file - self.port = port - self.client = None - - def connect(self): - """Establish SSH connection to the router""" - try: - self.client = paramiko.SSHClient() - self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - - if self.key_file: - self.client.connect( - hostname=self.hostname, - username=self.username, - key_filename=self.key_file, - port=self.port, - timeout=30 - ) - elif os.environ.get('SSH_AUTH_SOCK'): - # Use SSH agent if available - self.client.connect( - hostname=self.hostname, - username=self.username, - port=self.port, - timeout=30 - ) - else: - self.client.connect( - hostname=self.hostname, - username=self.username, - password=self.password, - port=self.port, - timeout=30 - ) - print(f"Successfully connected to {self.hostname}") - return True - except Exception as e: - print(f"Failed to connect to {self.hostname}: {e}") - return False - - def run_command(self, command): - """Execute a command on the router and return the output""" - if not self.client: - print("No active connection") - return None - - try: - stdin, stdout, stderr = self.client.exec_command(command) - output = stdout.read().decode('utf-8') - error = stderr.read().decode('utf-8') - - if error: - print(f"Error executing '{command}': {error}") - return None - - return output - except Exception as e: - print(f"Failed to execute command '{command}': {e}") - return None - - def backup_commands(self, commands, output_dir="backup"): - """Run multiple commands and save outputs to files""" - if not os.path.exists(output_dir): - os.makedirs(output_dir) - - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - backup_info = { - "hostname": self.hostname, - "timestamp": timestamp, - "commands": {} - } - - # Truncate output file at start - filename = f"{self.hostname}.txt" - filepath = os.path.join(output_dir, filename) - open(filepath, 'w').close() - - for i, command in enumerate(commands): - print(f"Running command {i+1}/{len(commands)}: {command}") - output = self.run_command(command) - - if output: - # Use hostname as filename - filename = f"{self.hostname}.txt" - filepath = os.path.join(output_dir, filename) - - with open(filepath, 'a') as f: - f.write(f"## COMMAND: {command}\n") - f.write(output) - - backup_info["commands"][command] = { - "filename": filename, - "success": True, - "output_length": len(output) - } - print(f"Output saved to {filepath}") - else: - backup_info["commands"][command] = { - "filename": None, - "success": False, - "output_length": 0 - } - - # Save backup info - info_file = os.path.join(output_dir, f"{self.hostname}_{timestamp}_info.json") - with open(info_file, 'w') as f: - json.dump(backup_info, f, indent=2) - - return backup_info - - def disconnect(self): - """Close SSH connection""" - if self.client: - self.client.close() - print(f"Disconnected from {self.hostname}") - - -def main(): - parser = argparse.ArgumentParser(description='SSH Router Backup Tool') - parser.add_argument('--config', required=True, help='YAML configuration file path') - parser.add_argument('--password', help='SSH password') - parser.add_argument('--key-file', help='SSH private key file path') - parser.add_argument('--port', type=int, default=22, help='SSH port (default: 22)') - parser.add_argument('--git-repo', default='/tmp', help='Git repository directory for command output files (default: /tmp)') - - args = parser.parse_args() - - # Load configuration - try: - with open(args.config, 'r') as f: - config = yaml.safe_load(f) - except Exception as e: - print(f"Failed to load config file {args.config}: {e}") - sys.exit(1) - - # Use SSH key by default, fall back to password - if not args.password and not args.key_file: - # Check if SSH agent is available first - if os.environ.get('SSH_AUTH_SOCK'): - print("Using SSH agent for authentication") - else: - # Try default SSH key locations - default_keys = [ - os.path.expanduser("~/.ssh/id_rsa"), - os.path.expanduser("~/.ssh/id_ed25519"), - os.path.expanduser("~/.ssh/id_ecdsa") - ] - - for key_path in default_keys: - if os.path.exists(key_path): - args.key_file = key_path - print(f"Using SSH key: {key_path}") - break - - # If no key found and no SSH agent, prompt for password - if not args.key_file: - import getpass - args.password = getpass.getpass("No SSH key found. Enter SSH password: ") - - # Process each device in config - devices = config.get('devices', {}) - if not devices: - print("No devices found in config file") - sys.exit(1) - - types = config.get('types', {}) - - success_count = 0 - total_count = len(devices) - - for hostname, device_config in devices.items(): - print(f"\nProcessing device: {hostname}") - - user = device_config.get('user') - commands = device_config.get('commands', []) - device_type = device_config.get('type') - - # If device has a type, get commands from types section - if device_type and device_type in types: - commands = types[device_type].get('commands', []) - - if not user: - print(f"No user specified for {hostname}, skipping") - continue - - if not commands: - print(f"No commands specified for {hostname}, skipping") - continue - - # Create backup instance - backup = RouterBackup( - hostname=hostname, - username=user, - password=args.password, - key_file=args.key_file, - port=args.port - ) - - # Connect and backup - if backup.connect(): - try: - backup_info = backup.backup_commands(commands, args.git_repo) - print(f"Backup completed for {hostname}") - print(f"Summary: {sum(1 for cmd in backup_info['commands'].values() if cmd['success'])}/{len(commands)} commands successful") - success_count += 1 - finally: - backup.disconnect() - else: - print(f"Failed to connect to {hostname}") - - print(f"\nOverall summary: {success_count}/{total_count} devices processed successfully") - - -if __name__ == "__main__": - main() diff --git a/src/go.mod b/src/go.mod new file mode 100644 index 0000000..37bde20 --- /dev/null +++ b/src/go.mod @@ -0,0 +1,15 @@ +module router_backup + +go 1.21 + +require ( + github.com/spf13/cobra v1.8.0 + golang.org/x/crypto v0.18.0 + gopkg.in/yaml.v2 v2.4.0 +) + +require ( + github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/spf13/pflag v1.0.5 // indirect + golang.org/x/sys v0.16.0 // indirect +) diff --git a/src/go.sum b/src/go.sum new file mode 100644 index 0000000..9dac6a3 --- /dev/null +++ b/src/go.sum @@ -0,0 +1,19 @@ +github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= +github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= +github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0= +github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho= +github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= +github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +golang.org/x/crypto v0.18.0 h1:PGVlW0xEltQnzFZ55hkuX5+KLyrMYhHld1YHO4AKcdc= +golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1mg= +golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU= +golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.16.0 h1:m+B6fahuftsE9qjo0VWp2FW0mB3MTJvR0BaMQrq0pmE= +golang.org/x/term v0.16.0/go.mod h1:yn7UURbUtPyrVJPGPq404EukNFxcm/foM+bV/bfcDsY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/src/router_backup.go b/src/router_backup.go new file mode 100644 index 0000000..08bd4fc --- /dev/null +++ b/src/router_backup.go @@ -0,0 +1,315 @@ +package main + +import ( + "fmt" + "io/ioutil" + "log" + "net" + "os" + "path/filepath" + "time" + + "github.com/spf13/cobra" + "golang.org/x/crypto/ssh" + "golang.org/x/crypto/ssh/agent" + "gopkg.in/yaml.v2" +) + +// Config structures +type Config struct { + Types map[string]DeviceType `yaml:"types"` + Devices map[string]Device `yaml:"devices"` +} + +type DeviceType struct { + Commands []string `yaml:"commands"` +} + +type Device struct { + User string `yaml:"user"` + Type string `yaml:"type,omitempty"` + Commands []string `yaml:"commands,omitempty"` +} + +// RouterBackup handles SSH connections and command execution +type RouterBackup struct { + hostname string + username string + password string + keyFile string + port int + client *ssh.Client +} + +// NewRouterBackup creates a new RouterBackup instance +func NewRouterBackup(hostname, username, password, keyFile string, port int) *RouterBackup { + return &RouterBackup{ + hostname: hostname, + username: username, + password: password, + keyFile: keyFile, + port: port, + } +} + +// Connect establishes SSH connection to the router +func (rb *RouterBackup) Connect() error { + config := &ssh.ClientConfig{ + User: rb.username, + HostKeyCallback: ssh.InsecureIgnoreHostKey(), + Timeout: 30 * time.Second, + } + + // Try SSH agent first if available + if sshAuthSock := os.Getenv("SSH_AUTH_SOCK"); sshAuthSock != "" { + if conn, err := net.Dial("unix", sshAuthSock); err == nil { + agentClient := agent.NewClient(conn) + config.Auth = []ssh.AuthMethod{ssh.PublicKeysCallback(agentClient.Signers)} + } + } + + // If SSH agent didn't work, try key file + if len(config.Auth) == 0 && rb.keyFile != "" { + key, err := ioutil.ReadFile(rb.keyFile) + if err != nil { + return fmt.Errorf("unable to read private key: %v", err) + } + + signer, err := ssh.ParsePrivateKey(key) + if err != nil { + return fmt.Errorf("unable to parse private key: %v", err) + } + + config.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)} + } + + // Fall back to password if available + if len(config.Auth) == 0 && rb.password != "" { + config.Auth = []ssh.AuthMethod{ssh.Password(rb.password)} + } + + if len(config.Auth) == 0 { + return fmt.Errorf("no authentication method available") + } + + address := fmt.Sprintf("%s:%d", rb.hostname, rb.port) + client, err := ssh.Dial("tcp", address, config) + if err != nil { + return fmt.Errorf("failed to connect to %s: %v", rb.hostname, err) + } + + rb.client = client + fmt.Printf("Successfully connected to %s\n", rb.hostname) + return nil +} + +// RunCommand executes a command on the router and returns the output +func (rb *RouterBackup) RunCommand(command string) (string, error) { + if rb.client == nil { + return "", fmt.Errorf("no active connection") + } + + session, err := rb.client.NewSession() + if err != nil { + return "", fmt.Errorf("failed to create session: %v", err) + } + defer session.Close() + + output, err := session.CombinedOutput(command) + if err != nil { + return "", fmt.Errorf("failed to execute command '%s': %v", command, err) + } + + return string(output), nil +} + +// BackupCommands runs multiple commands and saves outputs to files +func (rb *RouterBackup) BackupCommands(commands []string, gitRepo string) error { + if err := os.MkdirAll(gitRepo, 0755); err != nil { + return fmt.Errorf("failed to create directory %s: %v", gitRepo, err) + } + + filename := fmt.Sprintf("%s.txt", rb.hostname) + filepath := filepath.Join(gitRepo, filename) + + // Truncate file at start + file, err := os.Create(filepath) + if err != nil { + return fmt.Errorf("failed to create file %s: %v", filepath, err) + } + file.Close() + + successCount := 0 + for i, command := range commands { + fmt.Printf("Running command %d/%d: %s\n", i+1, len(commands), command) + output, err := rb.RunCommand(command) + + if err != nil { + fmt.Printf("Error executing '%s': %v\n", command, err) + continue + } + + // Append to file + file, err := os.OpenFile(filepath, os.O_APPEND|os.O_WRONLY, 0644) + if err != nil { + fmt.Printf("Failed to open file for writing: %v\n", err) + continue + } + + fmt.Fprintf(file, "## COMMAND: %s\n", command) + file.WriteString(output) + file.Close() + + fmt.Printf("Output saved to %s\n", filepath) + successCount++ + } + + fmt.Printf("Summary: %d/%d commands successful\n", successCount, len(commands)) + return nil +} + +// Disconnect closes SSH connection +func (rb *RouterBackup) Disconnect() { + if rb.client != nil { + rb.client.Close() + fmt.Printf("Disconnected from %s\n", rb.hostname) + } +} + +// loadConfig loads the YAML configuration file +func loadConfig(configPath string) (*Config, error) { + data, err := ioutil.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("failed to read config file %s: %v", configPath, err) + } + + var config Config + err = yaml.Unmarshal(data, &config) + if err != nil { + return nil, fmt.Errorf("failed to parse YAML: %v", err) + } + + return &config, nil +} + +// findDefaultSSHKey looks for default SSH keys +func findDefaultSSHKey() string { + homeDir, err := os.UserHomeDir() + if err != nil { + return "" + } + + defaultKeys := []string{ + filepath.Join(homeDir, ".ssh", "id_rsa"), + filepath.Join(homeDir, ".ssh", "id_ed25519"), + filepath.Join(homeDir, ".ssh", "id_ecdsa"), + } + + for _, keyPath := range defaultKeys { + if _, err := os.Stat(keyPath); err == nil { + fmt.Printf("Using SSH key: %s\n", keyPath) + return keyPath + } + } + + return "" +} + +func main() { + var configPath string + var password string + var keyFile string + var port int + var gitRepo string + + var rootCmd = &cobra.Command{ + Use: "router_backup", + Short: "SSH Router Backup Tool", + Long: "Connects to routers via SSH and runs commands, saving output to local files.", + Run: func(cmd *cobra.Command, args []string) { + // Load configuration + config, err := loadConfig(configPath) + if err != nil { + log.Fatalf("Failed to load config: %v", err) + } + + // Check authentication setup + if password == "" && keyFile == "" { + if os.Getenv("SSH_AUTH_SOCK") != "" { + fmt.Println("Using SSH agent for authentication") + } else { + keyFile = findDefaultSSHKey() + if keyFile == "" { + log.Fatal("No SSH key found and no password provided") + } + } + } + + // Process devices + if len(config.Devices) == 0 { + log.Fatal("No devices found in config file") + } + + successCount := 0 + totalCount := len(config.Devices) + + for hostname, deviceConfig := range config.Devices { + fmt.Printf("\nProcessing device: %s\n", hostname) + + user := deviceConfig.User + commands := deviceConfig.Commands + deviceType := deviceConfig.Type + + // If device has a type, get commands from types section + if deviceType != "" { + if typeConfig, exists := config.Types[deviceType]; exists { + commands = typeConfig.Commands + } + } + + if user == "" { + fmt.Printf("No user specified for %s, skipping\n", hostname) + continue + } + + if len(commands) == 0 { + fmt.Printf("No commands specified for %s, skipping\n", hostname) + continue + } + + // Create backup instance + backup := NewRouterBackup(hostname, user, password, keyFile, port) + + // Connect and backup + if err := backup.Connect(); err != nil { + fmt.Printf("Failed to connect to %s: %v\n", hostname, err) + continue + } + + err = backup.BackupCommands(commands, gitRepo) + backup.Disconnect() + + if err != nil { + fmt.Printf("Backup failed for %s: %v\n", hostname, err) + } else { + fmt.Printf("Backup completed for %s\n", hostname) + successCount++ + } + } + + fmt.Printf("\nOverall summary: %d/%d devices processed successfully\n", successCount, totalCount) + }, + } + + rootCmd.Flags().StringVar(&configPath, "config", "", "YAML configuration file path (required)") + 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(&gitRepo, "git-repo", "/tmp", "Git repository directory for command output files") + + rootCmd.MarkFlagRequired("config") + + if err := rootCmd.Execute(); err != nil { + log.Fatal(err) + } +} \ No newline at end of file