DebTik/cmd/debtik-cli/main.go

package main

import (
	"bufio"
	"encoding/json"
	"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")
	jsonOutput := flag.Bool("json", false, "Return a JSON envelope in single-command mode")
	passwordStdin := flag.Bool("password-stdin", false, "Read password from standard input (single-command mode)")
	timeout := flag.Duration("timeout", 60*time.Second, "Single-command timeout")
	cmdFlag := flag.String("c", "", "Execute a single command non-interactively and exit")
	flag.Parse()

	if *jsonOutput && *cmdFlag == "" {
		fmt.Fprintln(os.Stderr, "-json requires -c")
		os.Exit(1)
	}
	if *passwordStdin {
		if *cmdFlag == "" || *password != "" {
			fmt.Fprintln(os.Stderr, "-password-stdin requires -c and cannot be combined with -p")
			os.Exit(1)
		}
		line, err := bufio.NewReader(os.Stdin).ReadString('\n')
		if err != nil && !errors.Is(err, io.EOF) {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		*password = strings.TrimRight(line, "\r\n")
	}
	if *timeout <= 0 {
		fmt.Fprintln(os.Stderr, "-timeout must be positive")
		os.Exit(1)
	}
	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, *timeout, *jsonOutput); 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, timeout time.Duration, jsonOutput bool) 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, io.Discard); err != nil {
		return err
	}
	if strings.ContainsAny(command, "\r\n") {
		return fmt.Errorf("single-command mode accepts one line")
	}
	if err := conn.SetDeadline(time.Now().Add(timeout)); err != nil {
		return err
	}
	if _, err := fmt.Fprintf(conn, ":exec-status %s\n", command); err != nil {
		return err
	}
	data, err := io.ReadAll(reader)
	if err != nil {
		return fmt.Errorf("command failed or timed out: %w", err)
	}
	const begin = "DEBTIK-EXEC-BEGIN"
	const marker = "DEBTIK-EXEC-STATUS:"
	response := string(data)
	start := strings.Index(response, begin)
	index := strings.LastIndex(response, marker)
	if start < 0 || index < start+len(begin) {
		return fmt.Errorf("server did not return command status; update the server")
	}
	output := strings.TrimPrefix(response[start+len(begin):index], "\r\n")
	output = strings.TrimPrefix(output, "\n")
	status := strings.TrimSpace(response[index+len(marker):])
	if status != "0" && status != "1" {
		return fmt.Errorf("invalid command status")
	}
	if jsonOutput {
		if err := json.NewEncoder(os.Stdout).Encode(struct {
			Output string `json:"output"`
			Status int    `json:"status"`
			OK     bool   `json:"ok"`
		}{output, int(status[0] - '0'), status == "0"}); err != nil {
			return err
		}
	} else {
		fmt.Print(output)
	}
	if status != "0" {
		return fmt.Errorf("command failed (status %s)", status)
	}
	return nil
}

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, os.Stdout); 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, welcome io.Writer) 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(welcome, text)
			return err
		}
	}
	return fmt.Errorf("invalid authentication response")
}