main.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. package main
  2. import (
  3. "bufio"
  4. "errors"
  5. "flag"
  6. "fmt"
  7. "io"
  8. "net"
  9. "os"
  10. "os/signal"
  11. "strings"
  12. "syscall"
  13. "time"
  14. "golang.org/x/crypto/ssh/terminal"
  15. )
  16. func main() {
  17. addr := flag.String("addr", "127.0.0.1:8282", "DebTik CLI server address")
  18. username := flag.String("u", "", "Username for automatic authentication")
  19. password := flag.String("p", "", "Password for automatic authentication")
  20. cmdFlag := flag.String("c", "", "Execute a single command non-interactively and exit")
  21. flag.Parse()
  22. conn, err := net.DialTimeout("tcp", *addr, 3*time.Second)
  23. if err != nil {
  24. fmt.Fprintf(os.Stderr, "\033[31mError: Could not connect to DebTik CLI service on %s\033[0m\n", *addr)
  25. fmt.Fprintf(os.Stderr, "Ensure debtik service is active and listening on port 8282.\n")
  26. os.Exit(1)
  27. }
  28. defer conn.Close()
  29. // If one-shot command execution is requested:
  30. if *cmdFlag != "" {
  31. if err := runOneShot(conn, *username, *password, *cmdFlag); err != nil {
  32. fmt.Fprintln(os.Stderr, err)
  33. os.Exit(1)
  34. }
  35. return
  36. }
  37. // Interactive Mode
  38. if err := runInteractive(conn, *username, *password); err != nil {
  39. fmt.Fprintln(os.Stderr, err)
  40. os.Exit(1)
  41. }
  42. }
  43. func runOneShot(conn net.Conn, username, password, command string) error {
  44. reader := bufio.NewReader(conn)
  45. if username == "" || password == "" {
  46. return fmt.Errorf("single-command mode requires -u and -p")
  47. }
  48. if err := autoLogin(conn, reader, username, password); err != nil {
  49. return err
  50. }
  51. if _, err := fmt.Fprintf(conn, "%s\nquit\n", command); err != nil {
  52. return err
  53. }
  54. _, err := io.Copy(os.Stdout, reader)
  55. return err
  56. }
  57. func runInteractive(conn net.Conn, username, password string) error {
  58. reader := bufio.NewReader(conn)
  59. if username != "" || password != "" {
  60. if username == "" || password == "" {
  61. return fmt.Errorf("provide both -u and -p")
  62. }
  63. if err := autoLogin(conn, reader, username, password); err != nil {
  64. return err
  65. }
  66. }
  67. fd := int(os.Stdin.Fd())
  68. if terminal.IsTerminal(fd) {
  69. oldState, err := terminal.MakeRaw(fd)
  70. if err != nil {
  71. return err
  72. }
  73. defer terminal.Restore(fd, oldState)
  74. }
  75. sigChan := make(chan os.Signal, 1)
  76. signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
  77. defer signal.Stop(sigChan)
  78. errChan := make(chan error, 2)
  79. go func() { _, err := io.Copy(os.Stdout, reader); errChan <- err }()
  80. go func() { _, err := io.Copy(conn, os.Stdin); errChan <- err }()
  81. select {
  82. case <-sigChan:
  83. return nil
  84. case err := <-errChan:
  85. if errors.Is(err, io.EOF) {
  86. return nil
  87. }
  88. return err
  89. }
  90. }
  91. func autoLogin(conn net.Conn, reader *bufio.Reader, username, password string) error {
  92. if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
  93. return err
  94. }
  95. defer conn.SetDeadline(time.Time{})
  96. for _, prompt := range []struct{ text, reply string }{{"login:", username}, {"password:", password}} {
  97. var received strings.Builder
  98. for received.Len() < 65536 {
  99. ch, err := reader.ReadByte()
  100. if err != nil {
  101. return fmt.Errorf("authentication failed: %w", err)
  102. }
  103. received.WriteByte(ch)
  104. if strings.HasSuffix(strings.ToLower(received.String()), prompt.text) {
  105. if _, err := fmt.Fprintf(conn, "%s\n", prompt.reply); err != nil {
  106. return err
  107. }
  108. break
  109. }
  110. }
  111. if received.Len() >= 65536 {
  112. return fmt.Errorf("invalid login prompt")
  113. }
  114. }
  115. var received strings.Builder
  116. for received.Len() < 65536 {
  117. ch, err := reader.ReadByte()
  118. if err != nil {
  119. return fmt.Errorf("authentication failed: %w", err)
  120. }
  121. received.WriteByte(ch)
  122. text := received.String()
  123. if strings.Contains(strings.ToLower(text), "login failed:") {
  124. return fmt.Errorf("authentication rejected")
  125. }
  126. if strings.HasSuffix(text, "] > ") {
  127. _, err := fmt.Fprint(os.Stdout, text)
  128. return err
  129. }
  130. }
  131. return fmt.Errorf("invalid authentication response")
  132. }