// 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"
	"errors"
	"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
}

// dialSSH opens an authenticated SSH connection. The KEX list is widened
// because some RouterOS boxes (notably edge) only offer
// diffie-hellman-group-exchange-sha256, which x/crypto/ssh doesn't enable
// by default.
func dialSSH(r Router, sshPort int) (*ssh.Client, error) {
	addr := net.JoinHostPort(r.Host, strconv.Itoa(sshPort))

	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)
	}
	return client, nil
}

// 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) {
	client, err := dialSSH(r, sshPort)
	if err != nil {
		return nil, 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
}

// flattenRSC converts a .rsc-style script (with sticky section paths and
// backslash continuations) into a stream of self-contained ROS commands,
// one per line. ROS's non-interactive SSH shell treats each line
// independently — it does NOT remember the most recent section path the
// way /import or the interactive prompt do — so a script like
//
//	/ipv6 address
//	add address=... interface=...
//	add address=... interface=...
//
// fails with "bad command name add" unless we rewrite it as
//
//	/ipv6 address add address=... interface=...
//	/ipv6 address add address=... interface=...
//
// before piping it in. This function does that rewrite and also: drops
// `#` comments, drops blank lines, and joins lines ending in `\`.
func flattenRSC(script []byte) []byte {
	lines := strings.Split(string(script), "\n")

	// Collapse "\" line-continuations first so we can reason per logical line.
	var joined []string
	var carry strings.Builder
	for _, line := range lines {
		s := strings.TrimRight(line, "\r")
		if strings.HasSuffix(s, "\\") {
			carry.WriteString(strings.TrimSuffix(s, "\\"))
			carry.WriteString(" ")
			continue
		}
		if carry.Len() > 0 {
			carry.WriteString(s)
			joined = append(joined, carry.String())
			carry.Reset()
			continue
		}
		joined = append(joined, s)
	}
	if carry.Len() > 0 {
		joined = append(joined, carry.String())
	}

	var out strings.Builder
	section := ""
	for _, raw := range joined {
		s := strings.TrimSpace(raw)
		if s == "" || strings.HasPrefix(s, "#") {
			continue
		}
		if strings.HasPrefix(s, "/") {
			// A line starting with `/` is either a section header
			// (e.g. `/ipv6 address`) or an already-flat command
			// (e.g. `/ipv6 settings set forward=yes`). We detect a
			// flat command by the presence of a verb token after
			// the path words. Heuristic: if the line contains
			// ` add ` / ` set ` / ` remove ` / ` print ` / ` get`,
			// it's a flat command — emit as-is and don't change
			// the section. Otherwise treat the whole line as a
			// new section path.
			rest := s
			isFlat := false
			for _, verb := range []string{" add ", " add\t", " set ", " set\t", " remove ", " remove\t", " print", " get", " import "} {
				if strings.Contains(rest, verb) || strings.HasSuffix(rest, " add") || strings.HasSuffix(rest, " set") || strings.HasSuffix(rest, " remove") || strings.HasSuffix(rest, " print") || strings.HasSuffix(rest, " get") {
					isFlat = true
					break
				}
			}
			if isFlat {
				out.WriteString(s)
				out.WriteByte('\n')
				continue
			}
			section = s
			continue
		}
		if section == "" {
			// Defensive: command with no preceding section. Send
			// as-is and let RouterOS report the error.
			out.WriteString(s)
			out.WriteByte('\n')
			continue
		}
		out.WriteString(section)
		out.WriteByte(' ')
		out.WriteString(s)
		out.WriteByte('\n')
	}
	return []byte(out.String())
}

// runScriptOverSSH pipes a multi-line .rsc script through an SSH session as
// if the user had typed it. The script is first flattened (section paths
// inlined, comments stripped, `\` continuations joined) because ROS's
// non-interactive SSH shell forgets the section path between lines.
// Per-line errors print but don't abort the rest of the script. A trailing
// "/quit\n" is appended so the shell exits when the file ends.
func runScriptOverSSH(r Router, sshPort int, script []byte) ([]byte, error) {
	client, err := dialSSH(r, sshPort)
	if err != nil {
		return nil, 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

	body := flattenRSC(script)
	if len(body) == 0 || body[len(body)-1] != '\n' {
		body = append(body, '\n')
	}
	body = append(body, []byte("/quit\n")...)
	session.Stdin = bytes.NewReader(body)

	if err := session.Shell(); err != nil {
		return nil, fmt.Errorf("ssh shell: %w", err)
	}

	if err := session.Wait(); err != nil {
		// ROS shells exit with status 1 on /quit (they don't distinguish
		// "session ended normally" from "command failed"), and sometimes
		// close the channel without any exit status at all. Both look
		// like errors to ssh.Session.Wait(); treat them as success.
		// Real per-line failures show up in the captured stdout where
		// the user can see them.
		var exitErr *ssh.ExitError
		var exitMissing *ssh.ExitMissingError
		if !errors.As(err, &exitErr) && !errors.As(err, &exitMissing) {
			return stdout.Bytes(), fmt.Errorf("ssh wait: %w (stderr: %s)", 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()
}

// findRouter looks up a router by name in the config; returns nil if missing.
func findRouter(cfg *Config, name string) *Router {
	for i := range cfg.Routers {
		if cfg.Routers[i].Name == name {
			r := cfg.Routers[i]
			return &r
		}
	}
	return nil
}

// 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) {
	target := findRouter(cfg, routerName)
	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, "  "))
	}
}

// runSchema introspects a MikroTik API path on a single router and shows
// what parameter names exist in the wild. ROS7 keeps shuffling property
// names between point releases (address-families → afi, type=blackhole
// → blackhole flag, etc.), and the live router + the .rsc exports are
// the only authoritative reference. Two sources, in order:
//
//  1. Live API: union of property keys observed across all rows under
//     <path>/print, with one example value per key. Falls back to /get
//     for singletons (/system/identity, /ip/dns, etc.).
//
//  2. Local .rsc exports: any block whose first line is the matching
//     ROS-script path (e.g. `/routing bgp connection` for input
//     `/routing/bgp/connection`), with the next several non-blank lines
//     so the user can see actual `add ...` invocations.
//
// Trailing /add /set /remove /print /get on the input path are stripped
// — typing `mikrotik-tool schema verona /routing/bgp/connection/add` is
// the natural reflex when you've just hit "unknown parameter" on an add.
func runSchema(cfg *Config, routerName, path string) {
	target := findRouter(cfg, routerName)
	if target == nil {
		log.Fatalf("no router named %q in config", routerName)
	}
	r := cfg.resolve(*target)

	norm := strings.TrimPrefix(path, "/")
	norm = strings.Trim(norm, "/")
	for _, suffix := range []string{"/add", "/set", "/remove", "/print", "/get", "/getall"} {
		norm = strings.TrimSuffix(norm, suffix)
	}
	apiPath := "/" + norm

	fmt.Printf("=== schema: %s on %s ===\n", apiPath, r.Name)

	c, err := dialAPI(r)
	if err != nil {
		log.Fatalf("[%s] %v", r.Name, err)
	}
	defer c.Close()

	keys := map[string]string{}
	source := ""

	if rep, err := c.RunArgs([]string{apiPath + "/print"}); err == nil && len(rep.Re) > 0 {
		source = fmt.Sprintf("union of keys across %d existing item(s)", len(rep.Re))
		for _, row := range rep.Re {
			for k, v := range row.Map {
				if _, exists := keys[k]; !exists {
					keys[k] = v
				}
			}
		}
	} else if rep, err := c.RunArgs([]string{apiPath + "/get"}); err == nil && len(rep.Re) > 0 {
		source = "singleton (/get)"
		for _, row := range rep.Re {
			for k, v := range row.Map {
				keys[k] = v
			}
		}
	}

	if len(keys) == 0 {
		fmt.Printf("\n# live api: no rows visible at %s\n", apiPath)
	} else {
		fmt.Printf("\n# live api: %s\n", source)
		var names []string
		for k := range keys {
			names = append(names, k)
		}
		sort.Strings(names)
		for _, k := range names {
			example := keys[k]
			if len(example) > 60 {
				example = example[:60] + "…"
			}
			fmt.Printf("  %-32s  e.g. %s\n", k, example)
		}
	}

	rscPath := "/" + strings.ReplaceAll(norm, "/", " ")
	fmt.Printf("\n# .rsc examples (matching `%s`):\n", rscPath)
	rscFiles, _ := filepath.Glob("*.rsc")
	sort.Strings(rscFiles)
	found := 0
	const maxExamples = 6
	const maxLinesPerExample = 12
	for _, f := range rscFiles {
		if found >= maxExamples {
			break
		}
		data, err := os.ReadFile(f)
		if err != nil {
			continue
		}
		lines := strings.Split(string(data), "\n")
		for i, line := range lines {
			if found >= maxExamples {
				break
			}
			if strings.TrimSpace(line) != rscPath {
				continue
			}
			fmt.Printf("\n  [%s:%d]\n", f, i+1)
			for j := 1; j < maxLinesPerExample && i+j < len(lines); j++ {
				next := lines[i+j]
				trimmed := strings.TrimSpace(next)
				if strings.HasPrefix(trimmed, "/") {
					break
				}
				if trimmed == "" {
					break
				}
				fmt.Printf("    %s\n", strings.TrimRight(next, "\r"))
			}
			found++
		}
	}
	if found == 0 {
		fmt.Println("  (no matches in local .rsc files)")
	} else if found >= maxExamples {
		fmt.Printf("\n  (truncated at %d examples)\n", maxExamples)
	}
}

// runScript pushes a local .rsc file to one router and runs it line-by-line
// over SSH. Same login flow as `export` (API → discover SSH port → SSH).
// The file is fed to the RouterOS shell exactly as if the user pasted it
// in WinBox/SSH; section paths persist across lines, comments are ignored
// by RouterOS, and an error on one line doesn't abort the rest.
func runScript(cfg *Config, routerName, scriptPath string) {
	target := findRouter(cfg, routerName)
	if target == nil {
		log.Fatalf("no router named %q in config", routerName)
	}
	r := cfg.resolve(*target)

	body, err := os.ReadFile(scriptPath)
	if err != nil {
		log.Fatalf("read %s: %v", scriptPath, err)
	}

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

	log.Printf("[%s] ssh %s:%d running %s (%d bytes)", r.Name, r.Host, sshPort, scriptPath, len(body))
	out, err := runScriptOverSSH(r, sshPort, body)
	if len(out) > 0 {
		os.Stdout.Write(out)
		if !bytes.HasSuffix(out, []byte("\n")) {
			fmt.Println()
		}
	}
	if err != nil {
		log.Fatalf("[%s] %v", r.Name, err)
	}
}

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
  api       run an arbitrary API command on one router
            (e.g. mikrotik-tool api verona /ip/route/print)
  schema    show parameter names accepted under an API path on one router,
            plus matching examples from local .rsc exports
            (e.g. mikrotik-tool schema verona /routing/bgp/connection)
  script    push a local .rsc file to one router and run it via SSH; output
            from the router (errors, /export prints, etc.) is mirrored to
            stdout (e.g. mikrotik-tool script 494 494-v6-apply.rsc)`)
	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:])
	case "schema":
		if len(args) < 3 {
			fmt.Fprintln(os.Stderr, "usage: mikrotik-tool schema <router> <api-path>")
			os.Exit(2)
		}
		runSchema(cfg, args[1], args[2])
	case "script":
		if len(args) < 3 {
			fmt.Fprintln(os.Stderr, "usage: mikrotik-tool script <router> <file.rsc>")
			os.Exit(2)
		}
		runScript(cfg, args[1], args[2])
	default:
		fmt.Fprintf(os.Stderr, "unknown command: %s\n", args[0])
		usage()
	}
}