// mikrotik-tool: connect to MikroTik routers listed in a YAML config and
// run various operations against them.
//
// Usage:
//
//	mikrotik-tool [-c routers.yaml] <command> [args]
//
// Commands:
//
//	list     Print routers from the config.
//	export   For each router: log in over API-SSL, ask /ip/service for the SSH
//	         port, then SSH in, run "/export", and save stdout to <name>.rsc
//	         in the current working directory.
//	inventory  Walk every router, pull /ip/address, /ip/arp, /ip/pool, and
//	           /ip/neighbor, SNMP-probe each candidate access point and backhaul
//	           radio for sysName, and re-render inventory.yaml in place.
//	           Customer-side equipment (anything in a *-cpe / cgnat pool, or
//	           anything LLDP-tagged as wlan-ap / station-only) is excluded.
//	           Tower switches discovered via LLDP/CDP land under `switches:`.
package main

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-routeros/routeros/v3"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"gopkg.in/yaml.v3"
)

type Defaults struct {
	Username string `yaml:"username"`
	Password string `yaml:"password"`
	Port     int    `yaml:"port"`
}

type Router struct {
	Name     string `yaml:"name"`
	Host     string `yaml:"host"`
	Username string `yaml:"username,omitempty"`
	Password string `yaml:"password,omitempty"`
	Port     int    `yaml:"port,omitempty"`
}

type Config struct {
	Defaults Defaults `yaml:"defaults"`
	Routers  []Router `yaml:"routers"`
}

func loadConfig(path string) (*Config, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", path, err)
	}
	var c Config
	if err := yaml.Unmarshal(data, &c); err != nil {
		return nil, fmt.Errorf("parse %s: %w", path, err)
	}
	return &c, nil
}

// resolve fills in any router-level fields from the defaults block.
func (c *Config) resolve(r Router) Router {
	if r.Username == "" {
		r.Username = c.Defaults.Username
	}
	if r.Password == "" {
		r.Password = c.Defaults.Password
	}
	if r.Port == 0 {
		r.Port = c.Defaults.Port
	}
	if r.Port == 0 {
		r.Port = 8729 // MikroTik API-SSL default
	}
	return r
}

// dialAPI opens an API-SSL connection (port 8729) and logs in.
// MikroTik routers typically present a self-signed cert, so we skip verification.
func dialAPI(r Router) (*routeros.Client, error) {
	addr := net.JoinHostPort(r.Host, strconv.Itoa(r.Port))
	tlsCfg := &tls.Config{InsecureSkipVerify: true}

	conn, err := tls.DialWithDialer(&net.Dialer{Timeout: 15 * time.Second}, "tcp", addr, tlsCfg)
	if err != nil {
		return nil, fmt.Errorf("tls dial %s: %w", addr, err)
	}
	c, err := routeros.NewClient(conn)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("new client: %w", err)
	}
	if err := c.Login(r.Username, r.Password); err != nil {
		c.Close()
		return nil, fmt.Errorf("login: %w", err)
	}
	return c, nil
}

// lookupSSHPort asks /ip/service which port the ssh service is bound to and
// whether it's enabled. Returns the port, or an error if the service is
// disabled or absent.
func lookupSSHPort(c *routeros.Client) (int, error) {
	rep, err := c.RunArgs([]string{
		"/ip/service/print",
		"?name=ssh",
		"=.proplist=name,port,disabled",
	})
	if err != nil {
		return 0, fmt.Errorf("/ip/service/print: %w", err)
	}
	if len(rep.Re) == 0 {
		return 0, fmt.Errorf("ssh service not present in /ip/service")
	}
	row := rep.Re[0].Map
	if row["disabled"] == "true" {
		return 0, fmt.Errorf("ssh service is disabled on this router")
	}
	port, err := strconv.Atoi(row["port"])
	if err != nil {
		return 0, fmt.Errorf("ssh port %q is not numeric: %w", row["port"], err)
	}
	return port, nil
}

// sshAuthMethods builds the auth method list, in preference order:
//   - SSH agent (if SSH_AUTH_SOCK is set), so a router that has the user's
//     pubkey installed and password auth disabled still works
//   - on-disk identity files (~/.ssh/id_ed25519, id_rsa, id_ecdsa)
//   - password from the YAML config
//   - keyboard-interactive replaying the same password
func sshAuthMethods(password string) []ssh.AuthMethod {
	var methods []ssh.AuthMethod

	if sock := os.Getenv("SSH_AUTH_SOCK"); sock != "" {
		if conn, err := net.Dial("unix", sock); err == nil {
			ag := agent.NewClient(conn)
			// Only register the agent if it actually holds keys. An empty
			// agent still counts as a failed publickey attempt against
			// servers like RouterOS that limit retries, which then blocks
			// the on-disk key auth that follows.
			if signers, err := ag.Signers(); err == nil && len(signers) > 0 {
				methods = append(methods, ssh.PublicKeys(signers...))
			}
		}
	}

	if home, err := os.UserHomeDir(); err == nil {
		var signers []ssh.Signer
		for _, name := range []string{"id_ed25519", "id_rsa", "id_ecdsa"} {
			data, err := os.ReadFile(filepath.Join(home, ".ssh", name))
			if err != nil {
				continue
			}
			signer, err := ssh.ParsePrivateKey(data)
			if err != nil {
				continue
			}
			signers = append(signers, signer)
		}
		if len(signers) > 0 {
			methods = append(methods, ssh.PublicKeys(signers...))
		}
	}

	methods = append(methods, ssh.Password(password))
	methods = append(methods, ssh.KeyboardInteractive(func(_, _ string, questions []string, _ []bool) ([]string, error) {
		answers := make([]string, len(questions))
		for i := range questions {
			answers[i] = password
		}
		return answers, nil
	}))
	return methods
}

// runOverSSH opens an SSH connection to host:port, runs cmd non-interactively,
// and returns whatever the router wrote to stdout.
func runOverSSH(r Router, sshPort int, cmd string) ([]byte, error) {
	addr := net.JoinHostPort(r.Host, strconv.Itoa(sshPort))

	// Some RouterOS boxes (notably the edge router) only offer
	// diffie-hellman-group-exchange-sha256 for KEX, which x/crypto/ssh keeps
	// out of its defaults. Opt back in by appending it to the supported list.
	algos := ssh.SupportedAlgorithms()
	kex := append(algos.KeyExchanges, ssh.KeyExchangeDHGEXSHA256)

	var banner string
	cfg := &ssh.ClientConfig{
		Config:          ssh.Config{KeyExchanges: kex},
		User:            r.Username,
		Auth:            sshAuthMethods(r.Password),
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		BannerCallback: func(msg string) error {
			banner = msg
			return nil
		},
		Timeout: 15 * time.Second,
	}

	client, err := ssh.Dial("tcp", addr, cfg)
	if err != nil {
		if banner != "" {
			return nil, fmt.Errorf("ssh dial %s: %w (server banner: %q)", addr, err, banner)
		}
		return nil, fmt.Errorf("ssh dial %s: %w", addr, err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("ssh session: %w", err)
	}
	defer session.Close()

	var stdout, stderr bytes.Buffer
	session.Stdout = &stdout
	session.Stderr = &stderr

	if err := session.Run(cmd); err != nil {
		return nil, fmt.Errorf("ssh run %q: %w (stderr: %s)", cmd, err, stderr.String())
	}
	return stdout.Bytes(), nil
}

// exportRouter logs in over the API to discover the SSH port, then SSHes in
// and runs "/export", writing stdout to <name>.rsc locally.
func exportRouter(cfg *Config, r Router) error {
	r = cfg.resolve(r)
	if r.Host == "" || r.Name == "" {
		return fmt.Errorf("router needs both name and host")
	}

	log.Printf("[%s] api %s@%s:%d", r.Name, r.Username, r.Host, r.Port)
	api, err := dialAPI(r)
	if err != nil {
		return err
	}
	sshPort, err := lookupSSHPort(api)
	api.Close()
	if err != nil {
		return err
	}

	log.Printf("[%s] ssh %s:%d /export", r.Name, r.Host, sshPort)
	body, err := runOverSSH(r, sshPort, "/export")
	if err != nil {
		return err
	}
	if len(body) == 0 {
		return fmt.Errorf("/export returned no output")
	}

	out := r.Name + ".rsc"
	if err := os.WriteFile(out, body, 0o644); err != nil {
		return err
	}
	log.Printf("[%s] wrote %s (%d bytes)", r.Name, out, len(body))
	return nil
}

func runExport(cfg *Config) {
	var wg sync.WaitGroup
	for _, r := range cfg.Routers {
		wg.Add(1)
		go func(r Router) {
			defer wg.Done()
			if err := exportRouter(cfg, r); err != nil {
				log.Printf("[%s] ERROR: %v", r.Name, err)
			}
		}(r)
	}
	wg.Wait()
}

// runAPI logs into one router by name and runs an arbitrary API command,
// printing the rows returned. Useful for ad-hoc lookups (`arp culleoka`).
func runAPI(cfg *Config, routerName string, apiArgs []string) {
	var target *Router
	for i := range cfg.Routers {
		if cfg.Routers[i].Name == routerName {
			r := cfg.Routers[i]
			target = &r
			break
		}
	}
	if target == nil {
		log.Fatalf("no router named %q in config", routerName)
	}
	r := cfg.resolve(*target)
	c, err := dialAPI(r)
	if err != nil {
		log.Fatalf("[%s] %v", r.Name, err)
	}
	defer c.Close()

	rep, err := c.RunArgs(apiArgs)
	if err != nil {
		log.Fatalf("[%s] %v: %v", r.Name, apiArgs, err)
	}
	for _, row := range rep.Re {
		var keys []string
		for k := range row.Map {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		var parts []string
		for _, k := range keys {
			parts = append(parts, fmt.Sprintf("%s=%s", k, row.Map[k]))
		}
		fmt.Println(strings.Join(parts, "  "))
	}
}

func runList(cfg *Config) {
	for _, r := range cfg.Routers {
		r2 := cfg.resolve(r)
		fmt.Printf("%-20s %s:%d  user=%s\n", r2.Name, r2.Host, r2.Port, r2.Username)
	}
}

func usage() {
	fmt.Fprintln(os.Stderr, `usage: mikrotik-tool [-c routers.yaml] <command>

commands:
  list      show routers loaded from the config
  export    ssh to each router, run "/export", save stdout to <name>.rsc locally
  inventory discover access points, backhaul radios, and tower switches across
            every router (skipping customer equipment) and re-render
            inventory.yaml in place`)
	os.Exit(2)
}

func main() {
	log.SetFlags(log.Ltime)

	cfgPath := "routers.yaml"
	args := os.Args[1:]

	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "-c", "--config":
			if i+1 >= len(args) {
				usage()
			}
			cfgPath = args[i+1]
			args = append(args[:i], args[i+2:]...)
			i--
		case "-h", "--help":
			usage()
		}
	}

	if len(args) == 0 {
		usage()
	}

	cfg, err := loadConfig(cfgPath)
	if err != nil {
		log.Fatal(err)
	}
	if len(cfg.Routers) == 0 {
		log.Fatalf("no routers in %s", cfgPath)
	}

	switch args[0] {
	case "list":
		runList(cfg)
	case "export":
		runExport(cfg)
	case "inventory":
		runInventory(cfg)
	case "api":
		if len(args) < 3 {
			fmt.Fprintln(os.Stderr, "usage: mikrotik-tool api <router> <api-path> [k=v ...]")
			os.Exit(2)
		}
		runAPI(cfg, args[1], args[2:])
	default:
		fmt.Fprintf(os.Stderr, "unknown command: %s\n", args[0])
		usage()
	}
}
