DebTik/cmd/debtik-cli/main.go

package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh/terminal"
)

func main() {
	addr := flag.String("addr", "127.0.0.1:8282", "DebTik CLI server address")
	username := flag.String("u", "", "Username for automatic authentication")
	password := flag.String("p", "", "Password for automatic authentication")
	cmdFlag := flag.String("c", "", "Execute a single command non-interactively and exit")
	flag.Parse()

	conn, err := net.DialTimeout("tcp", *addr, 3*time.Second)
	if err != nil {
		fmt.Fprintf(os.Stderr, "\033[31mError: Could not connect to DebTik CLI service on %s\033[0m\n", *addr)
		fmt.Fprintf(os.Stderr, "Ensure debtik service is active and listening on port 8282.\n")
		os.Exit(1)
	}
	defer conn.Close()

	// If one-shot command execution is requested:
	if *cmdFlag != "" {
		if err := runOneShot(conn, *username, *password, *cmdFlag); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return
	}

	// Interactive Mode
	if err := runInteractive(conn, *username, *password); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func runOneShot(conn net.Conn, username, password, command string) error {
	reader := bufio.NewReader(conn)
	if username == "" || password == "" {
		return fmt.Errorf("single-command mode requires -u and -p")
	}
	if err := autoLogin(conn, reader, username, password); err != nil {
		return err
	}
	if _, err := fmt.Fprintf(conn, "%s\nquit\n", command); err != nil {
		return err
	}
	_, err := io.Copy(os.Stdout, reader)
	return err
}

func runInteractive(conn net.Conn, username, password string) error {
	reader := bufio.NewReader(conn)
	if username != "" || password != "" {
		if username == "" || password == "" {
			return fmt.Errorf("provide both -u and -p")
		}
		if err := autoLogin(conn, reader, username, password); err != nil {
			return err
		}
	}
	fd := int(os.Stdin.Fd())
	if terminal.IsTerminal(fd) {
		oldState, err := terminal.MakeRaw(fd)
		if err != nil {
			return err
		}
		defer terminal.Restore(fd, oldState)
	}
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	defer signal.Stop(sigChan)
	errChan := make(chan error, 2)
	go func() { _, err := io.Copy(os.Stdout, reader); errChan <- err }()
	go func() { _, err := io.Copy(conn, os.Stdin); errChan <- err }()
	select {
	case <-sigChan:
		return nil
	case err := <-errChan:
		if errors.Is(err, io.EOF) {
			return nil
		}
		return err
	}
}

func autoLogin(conn net.Conn, reader *bufio.Reader, username, password string) error {
	if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
		return err
	}
	defer conn.SetDeadline(time.Time{})
	for _, prompt := range []struct{ text, reply string }{{"login:", username}, {"password:", password}} {
		var received strings.Builder
		for received.Len() < 65536 {
			ch, err := reader.ReadByte()
			if err != nil {
				return fmt.Errorf("authentication failed: %w", err)
			}
			received.WriteByte(ch)
			if strings.HasSuffix(strings.ToLower(received.String()), prompt.text) {
				if _, err := fmt.Fprintf(conn, "%s\n", prompt.reply); err != nil {
					return err
				}
				break
			}
		}
		if received.Len() >= 65536 {
			return fmt.Errorf("invalid login prompt")
		}
	}
	var received strings.Builder
	for received.Len() < 65536 {
		ch, err := reader.ReadByte()
		if err != nil {
			return fmt.Errorf("authentication failed: %w", err)
		}
		received.WriteByte(ch)
		text := received.String()
		if strings.Contains(strings.ToLower(text), "login failed:") {
			return fmt.Errorf("authentication rejected")
		}
		if strings.HasSuffix(text, "] > ") {
			_, err := fmt.Fprint(os.Stdout, text)
			return err
		}
	}
	return fmt.Errorf("invalid authentication response")
}