main.go 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273
  1. package main
  2. import (
  3. "bufio"
  4. "errors"
  5. "flag"
  6. "fmt"
  7. "io/ioutil"
  8. "os"
  9. "os/exec"
  10. "strings"
  11. "github.com/apcera/termtables"
  12. "github.com/ttacon/chalk"
  13. )
  14. /* TODO:
  15. * Detect non-terminal output and remove color. https://github.com/ttacon/chalk/issues/4
  16. * More correctly support ssh options like port: user@port:22
  17. * Parse and filter ssh -v to see when we are connected for slow logins.
  18. * Recent history, maybe even rain last to login to last server.
  19. * Configuration to disable unwanted features.
  20. * Adding does not detect when it is overwriting an existing server.
  21. * Smarter detection for typos in friendly name.
  22. * Add a flag for verbose/debug output
  23. * Add a flag to prevent a stored auto run command from executing
  24. */
  25. func main() {
  26. flag.Usage = usage
  27. //chalk.DetectTerminal()
  28. parseArgs()
  29. }
  30. func usage() {
  31. fmt.Println("☔ ./rain <command> [options]")
  32. fmt.Println()
  33. fmt.Println("Commands:")
  34. fmt.Println(" list")
  35. fmt.Println(" ssh <alias> [command(s)]")
  36. fmt.Println(" add [alias] [root@][hostname][:22] [command(s)]")
  37. fmt.Println(" note <alias>")
  38. fmt.Println(" search <alias|hostname|notes>")
  39. fmt.Println(" delete <alias>")
  40. fmt.Println(" help")
  41. fmt.Println()
  42. fmt.Println("Report bugs at http://github.com/trashcan/rain/issues.")
  43. }
  44. func handleError(m error) {
  45. if m != nil {
  46. fmt.Printf("☔\t%s%s%s\n", chalk.Red, m.Error(), chalk.Reset)
  47. os.Exit(1)
  48. }
  49. }
  50. func handleWarning(m error) {
  51. if m != nil {
  52. fmt.Printf("☔\t%s%s%s\n", chalk.Yellow, m.Error(), chalk.Reset)
  53. }
  54. }
  55. func handleStatus(m string) {
  56. fmt.Printf("☔\t%s%s%s\n", chalk.Green, m, chalk.Reset)
  57. }
  58. func handleDebug(m string) {
  59. fmt.Printf("☔\t%s%s%s\n", chalk.Green, m, chalk.Reset)
  60. }
  61. func parseArgs() {
  62. requireArgs(os.Args[0], 1)
  63. args := os.Args[1:]
  64. switch args[0] {
  65. case "ssh":
  66. requireArgs("ssh", 2)
  67. cmdSSH(args[1], strings.Join(args[2:], " "))
  68. case "list":
  69. cmdList()
  70. case "add":
  71. cmdAdd()
  72. case "delete":
  73. requireArgs("delete", 2)
  74. cmdDelete(args[1])
  75. case "search":
  76. requireArgs("search", 2)
  77. cmdSearch(args[1])
  78. case "note":
  79. requireArgs("note", 2)
  80. cmdNote(args[1])
  81. case "edit":
  82. requireArgs("note", 2)
  83. cmdNote(args[1])
  84. case "help":
  85. flag.Usage()
  86. default:
  87. flag.Usage()
  88. handleError(fmt.Errorf("unknown subcommand: %s", os.Args[1]))
  89. }
  90. }
  91. func requireArgs(cmd string, count int) {
  92. args := os.Args[1:]
  93. if len(args) < count {
  94. flag.Usage()
  95. handleError(fmt.Errorf("%s requires more argument(s).\n", cmd))
  96. }
  97. }
  98. func cmdAdd() {
  99. var alias, hostname, runcmd string
  100. if len(os.Args) >= 4 {
  101. alias = os.Args[2]
  102. hostname = os.Args[3]
  103. runcmd = ""
  104. } else {
  105. scanner := bufio.NewScanner(os.Stdin)
  106. fmt.Print("Alias: ")
  107. scanner.Scan()
  108. alias = scanner.Text()
  109. fmt.Print("Hostname ([user]@<hostname>[:port]): ")
  110. scanner.Scan()
  111. hostname = scanner.Text()
  112. runcmd = ""
  113. }
  114. if len(os.Args) > 4 {
  115. runcmd = strings.Join(os.Args[4:], " ")
  116. }
  117. newServer := Server{
  118. Alias: alias,
  119. Hostname: hostname,
  120. Notes: string(""),
  121. RunCmd: runcmd,
  122. }
  123. dbw := DBWrapper{}
  124. err := dbw.AddServer(newServer)
  125. handleError(err)
  126. handleStatus(hostname + " added successfully.")
  127. }
  128. func cmdList() {
  129. dbw := DBWrapper{}
  130. servers, err := dbw.AllServers()
  131. handleError(err)
  132. renderServers(servers, "")
  133. }
  134. func cmdDelete(alias string) {
  135. dbw := DBWrapper{}
  136. err := dbw.DeleteServer(alias)
  137. handleError(err)
  138. }
  139. func cmdSearch(search string) {
  140. dbw := DBWrapper{}
  141. servers, err := dbw.ServerSearch(search)
  142. handleError(err)
  143. renderServers(servers, search)
  144. }
  145. func cmdNote(alias string) {
  146. dbw := DBWrapper{}
  147. s, err := dbw.GetServer(alias)
  148. handleError(err)
  149. newNote := openEditor(s.Notes)
  150. if s.Notes != string(newNote) {
  151. s.Notes = newNote
  152. err = dbw.UpdateServer(s)
  153. handleError(err)
  154. }
  155. }
  156. func openEditor(notes string) (newNote string) {
  157. // Create a tempfile
  158. file, err := ioutil.TempFile(os.TempDir(), "rain")
  159. handleError(err)
  160. // Delete it when done.
  161. defer os.Remove(file.Name())
  162. // Write the current notes into the file
  163. err = ioutil.WriteFile(file.Name(), []byte(notes), 0644)
  164. handleError(err)
  165. // Launch vim to edit the notes
  166. cmd := exec.Command("vim", []string{file.Name()}...)
  167. cmd.Stderr = os.Stderr
  168. cmd.Stdout = os.Stdout
  169. cmd.Stdin = os.Stdin
  170. err = cmd.Start()
  171. handleError(err)
  172. err = cmd.Wait()
  173. handleError(err)
  174. // Read the updated notes
  175. newNoteByte, err := ioutil.ReadFile(file.Name())
  176. handleError(err)
  177. return string(newNoteByte)
  178. }
  179. func cmdSSH(alias string, runcmd string) {
  180. dbw := DBWrapper{}
  181. s, err := dbw.GetServer(alias)
  182. if err != nil {
  183. search, _ := dbw.ServerSearch(alias)
  184. if len(search) == 0 {
  185. // Just ssh to the provided hostname.
  186. handleWarning(err)
  187. s = Server{Hostname: alias}
  188. } else if len(search) == 1 {
  189. // If there's one search result, ssh to it.
  190. handleWarning(errors.New("Matched one result, going to " + search[0].Hostname + "."))
  191. s = search[0]
  192. s.Hit++
  193. dbw.UpdateServer(s)
  194. } else {
  195. // Otherwise, list the search results and quit.
  196. renderServers(search, alias)
  197. return
  198. }
  199. } else {
  200. handleStatus(fmt.Sprintf("Connecting to %s.", s.Hostname))
  201. s.Hit++
  202. dbw.UpdateServer(s)
  203. }
  204. if len(runcmd) > 0 {
  205. // A command was passed in to run on remote server, pass it through to ssh
  206. s.RunCmd = runcmd
  207. }
  208. s.ssh()
  209. }
  210. func renderServers(servers []Server, highlight string) {
  211. if len(servers) == 0 {
  212. handleError(errors.New("No servers found."))
  213. }
  214. var ts = &termtables.TableStyle{
  215. SkipBorder: true,
  216. BorderX: "", BorderY: "", BorderI: "",
  217. PaddingLeft: 0, PaddingRight: 4,
  218. Width: 80,
  219. Alignment: termtables.AlignLeft,
  220. }
  221. t := termtables.CreateTable()
  222. t.Style = ts
  223. cb := chalk.Bold.TextStyle
  224. // TODO FIXME: These are adding a blank line above the headers.
  225. // t.AddHeaders("Alias", "Hostname", "Hits")
  226. t.AddRow(cb("Alias"), cb("Hostname"), cb("AutoRun Cmd"), cb("Hits"))
  227. for _, s := range servers {
  228. if highlight != "" {
  229. s.Alias = strings.Replace(s.Alias, highlight, fmt.Sprintf("%s%s%s", chalk.Green, highlight, chalk.Reset), 1)
  230. s.Hostname = strings.Replace(s.Hostname, highlight, fmt.Sprintf("%s%s%s", chalk.Green, highlight, chalk.Reset), 1)
  231. }
  232. t.AddRow(s.Alias, s.Hostname, s.RunCmd, s.Hit)
  233. }
  234. fmt.Printf(t.Render())
  235. }
  236. func renderNotes(s Server) {
  237. handleStatus(fmt.Sprintf("Notes for %s:", s.Alias))
  238. fmt.Println(s.Notes)
  239. }