Netpulse_SasS/agent/internal/ncmx/session.go

// Package ncmx — знімання конфігу з обладнання по CLI.
//
// Головна складність тут не в мережі, а в тому, що консоль мережевого
// пристрою — це не програмний інтерфейс. Немає ані коду завершення, ані
// довжини відповіді: єдиний спосіб зрозуміти, що команда відпрацювала —
// побачити знову запрошення. Тому весь модуль побудований навколо
// пошуку промпту в потоці байтів.
package ncmx

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strings"
	"time"
)

// ErrPromptTimeout — промпт не з'явився у відведений час.
//
// Окрема помилка, бо це найчастіша причина невдачі й лікується вона не
// повтором, а виправленням prompt_regex у профілі. Плутати її з
// мережевою відмовою означало б ганяти ретраї там, де вони не
// допоможуть.
var ErrPromptTimeout = errors.New("не дочекались запрошення командного рядка")

// ErrTooLarge — пристрій віддав більше, ніж дозволено.
var ErrTooLarge = errors.New("вивід перевищив ліміт")

// Conn — те, що вміє двонаправлений обмін символами. SSH-сесія й
// Telnet-з'єднання зводяться до нього, і решта коду про різницю не знає.
type Conn interface {
	io.Reader
	io.Writer
	Close() error
}

// Options — налаштування знімання.
type Options struct {
	// PromptRe — за чим упізнаємо, що пристрій готовий приймати команду.
	PromptRe *regexp.Regexp
	// MoreRe — посторінковий вивід. Якщо трапиться, шлемо MoreSend.
	MoreRe   *regexp.Regexp
	MoreSend string

	// MaxBytes — стеля на команду. Пристрій із зациклленим виводом
	// інакше з'їсть пам'ять агента, а не свою.
	MaxBytes int
	// IdleTimeout — скільки чекати наступного байта.
	IdleTimeout time.Duration

	// Transcript накопичує весь обмін для діагностики prompt_regex.
	Transcript *bytes.Buffer
}

// Стандартне «ще» для більшості вендорів. Пробіл, а не Enter: Enter на
// частині платформ гортає рядок, а не сторінку, і збір конфігу
// розтягується на тисячі проходів.
const DefaultMoreSend = " "

// DefaultMoreRe ловить типові запрошення посторінкового виводу.
var DefaultMoreRe = regexp.MustCompile(`(?i)--\s*more\s*--|<--- more --->|\x1b\[7m--More--`)

// CLI — жива сесія з пристроєм.
type CLI struct {
	conn Conn
	opt  Options

	// buf тримає прочитане, але ще не розібране. Промпт може прийти
	// розірваним між двома read(), тому склеювати доводиться тут.
	buf bytes.Buffer
}

func NewCLI(conn Conn, opt Options) *CLI {
	if opt.MaxBytes <= 0 {
		opt.MaxBytes = 16 << 20
	}
	if opt.IdleTimeout <= 0 {
		opt.IdleTimeout = 20 * time.Second
	}
	if opt.MoreRe == nil {
		opt.MoreRe = DefaultMoreRe
	}
	if opt.MoreSend == "" {
		opt.MoreSend = DefaultMoreSend
	}
	return &CLI{conn: conn, opt: opt}
}

// WaitPrompt читає до першого запрошення. Викликається одразу після
// підключення: більшість пристроїв вітаються банером, і виконувати
// команду до того, як він добіжить, — це отримати її відлуння всередині
// банера.
func (c *CLI) WaitPrompt(ctx context.Context) error {
	_, err := c.readUntilPrompt(ctx)
	return err
}

// Run виконує команду й повертає її вивід без відлуння та без
// фінального запрошення.
func (c *CLI) Run(ctx context.Context, cmd string) (string, error) {
	if _, err := c.conn.Write([]byte(cmd + "\n")); err != nil {
		return "", fmt.Errorf("надсилання %q: %w", cmd, err)
	}

	raw, err := c.readUntilPrompt(ctx)
	if err != nil {
		return raw, err
	}
	return cleanOutput(raw, cmd), nil
}

// readUntilPrompt читає, доки не побачить запрошення.
func (c *CLI) readUntilPrompt(ctx context.Context) (string, error) {
	chunk := make([]byte, 8192)

	for {
		if err := ctx.Err(); err != nil {
			return c.buf.String(), err
		}

		// Промпт шукаємо лише в хвості: конфіг може містити рядок, що
		// схожий на запрошення (banner motd, коментар), і пошук по
		// всьому накопиченому обривав би збір на середині.
		if loc := c.matchTail(c.opt.PromptRe); loc != nil {
			// Ріжемо ВЕСЬ рядок із запрошенням, а не від місця збігу.
			// Типовий шаблон «[>#]\s*$» збігається лише з символом
			// запрошення, тому обрізання за loc[0] лишало б у конфізі
			// ім'я пристрою — «…interface Gi0/1» і окремим рядком «sw1».
			out := c.buf.String()[:lineStart(c.buf.String(), loc[0])]
			c.buf.Reset()
			return out, nil
		}

		// Посторінковий вивід: відповідаємо й прибираємо саме
		// запрошення з тексту, щоб воно не потрапило в конфіг.
		if loc := c.matchTail(c.opt.MoreRe); loc != nil {
			s := c.buf.String()
			c.buf.Reset()
			// Запрошення пейджера часто стоїть у власному рядку разом
			// із керуючими послідовностями очищення — прибираємо рядок
			// цілком, інакше в конфізі лишаються пробіли й ESC[K.
			c.buf.WriteString(s[:lineStart(s, loc[0])])
			if _, err := c.conn.Write([]byte(c.opt.MoreSend)); err != nil {
				return c.buf.String(), fmt.Errorf("відповідь на --More--: %w", err)
			}
			continue
		}

		if c.buf.Len() > c.opt.MaxBytes {
			return c.buf.String(), ErrTooLarge
		}

		n, err := c.readWithDeadline(ctx, chunk)
		if n > 0 {
			c.buf.Write(chunk[:n])
			if c.opt.Transcript != nil {
				c.opt.Transcript.Write(chunk[:n])
			}
		}
		if err != nil {
			if errors.Is(err, errIdle) {
				return c.buf.String(), ErrPromptTimeout
			}
			if errors.Is(err, io.EOF) {
				// Пристрій закрив сесію. Якщо промпт так і не
				// з'явився — це не успіх, навіть коли текст схожий
				// на конфіг.
				return c.buf.String(), ErrPromptTimeout
			}
			return c.buf.String(), err
		}
	}
}

// matchTail шукає збіг лише в кінці накопиченого.
func (c *CLI) matchTail(re *regexp.Regexp) []int {
	if re == nil {
		return nil
	}
	s := c.buf.String()
	// Хвоста в 512 байтів вистачає будь-якому запрошенню й відрізає
	// збіги всередині тіла конфігу.
	const tail = 512
	off := 0
	if len(s) > tail {
		off = len(s) - tail
	}
	loc := re.FindStringIndex(s[off:])
	if loc == nil {
		return nil
	}
	return []int{off + loc[0], off + loc[1]}
}

// lineStart повертає позицію початку рядка, у якому стоїть pos.
func lineStart(s string, pos int) int {
	if pos <= 0 {
		return 0
	}
	if i := strings.LastIndexByte(s[:pos], byte('\n')); i >= 0 {
		return i + 1
	}
	return 0
}

var errIdle = errors.New("тиша в каналі")

// readWithDeadline читає з обмеженням на паузу між байтами.
//
// Саме на паузу, а не на всю операцію: збір конфігу великого шасі
// триває хвилини, і загальний дедлайн довелося б ставити навмання.
// Тиша ж означає одне з двох — пристрій завис або промпт не той.
func (c *CLI) readWithDeadline(ctx context.Context, p []byte) (int, error) {
	type res struct {
		n   int
		err error
	}
	ch := make(chan res, 1)
	go func() {
		n, err := c.conn.Read(p)
		ch <- res{n, err}
	}()

	t := time.NewTimer(c.opt.IdleTimeout)
	defer t.Stop()

	select {
	case r := <-ch:
		return r.n, r.err
	case <-t.C:
		return 0, errIdle
	case <-ctx.Done():
		return 0, ctx.Err()
	}
}

// ansiRe ловить керуючі послідовності терміналу.
//
// CSI (кольори, позиціювання, bracketed paste) і OSC (заголовок вікна).
// Без цього в збережений конфіг потрапляє те, що бачить око оператора,
// а не те, що налаштовано на пристрої: перший же diff показує зміну
// там, де змінився лише режим терміналу.
var ansiRe = regexp.MustCompile("\x1b\\[[0-9;?]*[ -/]*[@-~]" +
	"|\x1b\\][^\x07\x1b]*(?:\x07|\x1b\\\\)" +
	"|\x1b[=>NOc]")

// crRe згортає підряд ідучі \r перед переводом рядка.
//
// Псевдотермінал видає «\r\r\n» частіше, ніж здається: сам пристрій
// шле «\r\n», а PTY додає свій «\r». Без згортання кожен такий рядок
// подвоюється, і конфіг на 300 байтів приїжджає на 12 рядків замість 10.
var crRe = regexp.MustCompile(`\r+\n`)

// cleanOutput прибирає відлуння команди, керуючі коди й нормалізує
// переводи рядка.
//
// Відлуння прибирається саме порівнянням із надісланою командою, а не
// «викинути перший рядок»: частина пристроїв відлуння не робить, і
// сліпе відкидання з'їдало б перший рядок конфігу.
func cleanOutput(raw, cmd string) string {
	s := ansiRe.ReplaceAllString(raw, "")
	s = crRe.ReplaceAllString(s, "\n")
	s = strings.ReplaceAll(s, "\r", "\n")

	lines := strings.Split(s, "\n")

	// Відлуння шукаємо серед перших кількох рядків і за суфіксом.
	//
	// Суфіксом, бо в PTY команда часто друкується в тому ж рядку, що й
	// запрошення: «claude@host:~$ cat /etc/os-release». Порівняння
	// цілого рядка тут не спрацювало б, і запрошення поїхало б у конфіг.
	want := strings.TrimSpace(cmd)
	for i := 0; i < len(lines) && i < 3; i++ {
		l := strings.TrimSpace(lines[i])
		if l == want || (want != "" && strings.HasSuffix(l, want)) {
			lines = lines[i+1:]
			break
		}
	}

	// Хвостові порожні рядки — наслідок того, що ми відрізали промпт.
	for len(lines) > 0 && strings.TrimSpace(lines[len(lines)-1]) == "" {
		lines = lines[:len(lines)-1]
	}
	// Провідні порожні — наслідок того, що ми відрізали відлуння.
	for len(lines) > 0 && strings.TrimSpace(lines[0]) == "" {
		lines = lines[1:]
	}
	return strings.Join(lines, "\n")
}

// StripFirstLines відкидає службову шапку.
//
// Скільки саме — знає профіль: у Cisco IOS це три рядки «Building
// configuration…», у Eltex MES5448 — десять.
func StripFirstLines(s string, n int) string {
	if n <= 0 {
		return s
	}
	lines := strings.Split(s, "\n")
	if n >= len(lines) {
		return ""
	}
	return strings.Join(lines[n:], "\n")
}

func (c *CLI) Close() error { return c.conn.Close() }