Go SSH CLI Utility

Go lang is an open-source programming language that was originally developed by Google. We will look into a short Go script that stores SSH commands and executes them.

8 months ago

Latest Post Cloud Development Kit for Kubernetes (cdk8s) by Tyler Moon

Go lang is an open-source programming language that was originally developed by Google. With a growing community and library catalog Go has become a great language for lower-level system programming.

One of the great uses for Go is in writing Linux command-line utilities. In this short article, we will look into a short Go script that stores SSH commands and executes them. This utility is useful when you have many different servers to login to and don't want to remember all of the ssh commands needed.

Prerequisites

Setup

If you are not familiar with the Go syntax checkout the Go by Example website for a great reference.

Create a new file called tssh.go in a new directory. Add the following code to the new file:

// tssh.go

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"syscall"

	"github.com/urfave/cli" // Library that adds useful utilities for handling command line arguments
)

// Struct for managing data
type sshTemplate struct {
	Name    string
	Command string
}

// Utility for handling errors
func check(e error) {
	if e != nil {
		panic(e)
	}
}

// Save a new template to a json file
func writeToFile(data map[string]sshTemplate, filePath string) {
	file, _ := json.MarshalIndent(data, "", " ")
	err := ioutil.WriteFile(filePath, file, 0644)
	check(err)
}

// Read from the json and parse a map of sshTemplate structs
func readFromFile(filePath string) map[string]sshTemplate {
	dat, err := ioutil.ReadFile(filePath)
	check(err)

	res := map[string]sshTemplate{}
	json.Unmarshal(dat, &res)

	return res
}

// Add a new sshTemplate struct and save it to the file
func addTemplate(name string, command string, filePath string) *sshTemplate {
	templateArray := readFromFile(filePath)

	newTemplate := sshTemplate{
		Name:    name,
		Command: command,
	}

	templateArray[name] = newTemplate

	writeToFile(templateArray, filePath)

	fmt.Println(fmt.Sprintf("New template '%s' added!", newTemplate.Name))
	return &newTemplate
}

// Remove a template from the json file
func removeTemplate(name string, filePath string) *sshTemplate {
	templateArray := readFromFile(filePath)

	template := templateArray[name]

	delete(templateArray, name)

	writeToFile(templateArray, filePath)

	fmt.Println(fmt.Sprintf("Template '%s' removed!", name))
	return &template
}

// Execute an ssh command based on saved data
func executeCommand(template sshTemplate) {
	fmt.Println("Executing SSH command for", template.Command)

	binary, lookErr := exec.LookPath("ssh")
	check(lookErr)

	args := []string{"ssh", template.Command}
	env := os.Environ()

	execErr := syscall.Exec(binary, args, env)
	check(execErr)
}

func main() {
	app := &cli.App{
		// With no other arguments assume that the input is a Name of a saved template
		// so load it from the file and execute the ssh command
		Action: func(c *cli.Context) error {
			argName := c.Args().First()

			a := readFromFile("tempData.json")
			template := a[argName]
			executeCommand(template)
			fmt.Println(template.Command)
			return nil
		},
		Commands: []*cli.Command{
			{
				Name:    "template",
				Aliases: []string{"t"},
				Usage:   "options for task templates",
				Subcommands: []*cli.Command{
					// tssh template add <name> <command>
					{
						Name:  "add",
						Usage: "add a new template",
						Action: func(c *cli.Context) error {
							name := c.Args().First()
							command := c.Args().Get(1)

							addTemplate(name, command, "tempData.json")

							return nil
						},
					},
					// tssh template remove <name> <command>
					{
						Name:  "remove",
						Usage: "remove an existing template",
						Action: func(c *cli.Context) error {
							name := c.Args().First()

							removeTemplate(name, "tempData.json")
							return nil
						},
					},
				},
			},
		},
	}

	// Run the application
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

Running go run tssh.go template add test 192.168.0.1 will add a new ssh template called "test" that will execute the command ssh 192.168.0.1 when the go run tssh.go test command is run. Then using go run tssh.go template remove test will delete it.

Tests

Go has a relatively simple unit testing framework for testing programs. Create a new file called tssh_test.go and add the following code:

package main

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"reflect"
	"testing"
)

func Test_addTemplate(t *testing.T) {
	filePath := "test.json"
	os.Create(filePath)

	type args struct {
		name    string
		command string
	}
	tests := []struct {
		name string
		args args
		want *sshTemplate
	}{
		{"testName", args{name: "testName", command: "testCommand"}, &sshTemplate{Name: "testName", Command: "testCommand"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := addTemplate(tt.args.name, tt.args.command, filePath); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("addTemplate() = %v, want %v", got, tt.want)
			}
		})
	}

	os.Remove(filePath)
}

func Test_removeTemplate(t *testing.T) {
	data := map[string]sshTemplate{"testName": {Name: "testName", Command: "testCommand"}}
	file, _ := json.MarshalIndent(data, "", " ")
	filePath := "test.json"
	os.Create(filePath)

	ioutil.WriteFile(filePath, file, 0644)

	type args struct {
		name string
	}
	tests := []struct {
		name string
		args args
		want *sshTemplate
	}{
		{"testName", args{name: "testName"}, &sshTemplate{Name: "testName", Command: "testCommand"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := removeTemplate(tt.args.name, filePath); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("removeTemplate() = %v, want %v", got, tt.want)
			}
		})
	}

	os.Remove(filePath)
}

func Test_readFromFile(t *testing.T) {
	data := map[string]sshTemplate{"testName": {Name: "testName", Command: "testCommand"}}
	file, _ := json.MarshalIndent(data, "", " ")
	filePath := "test.json"
	os.Create(filePath)

	ioutil.WriteFile(filePath, file, 0644)

	t.Run("test reading", func(t *testing.T) {
		if got := readFromFile(filePath); !reflect.DeepEqual(got, data) {
			t.Errorf("readFromFile() = %v, want %v", got, data)
		}
	})

	os.Remove(filePath)
}

func Test_writeToFile(t *testing.T) {
	data := map[string]sshTemplate{"testName": {Name: "testName", Command: "testCommand"}}
	filePath := "test.json"
	os.Create(filePath)

	t.Run("test writing", func(t *testing.T) {
		writeToFile(data, filePath)

		fileData, _ := ioutil.ReadFile(filePath)
		res := map[string]sshTemplate{}
		json.Unmarshal(fileData, &res)

		if !reflect.DeepEqual(res, data) {
			t.Errorf("writeToFile() = %v, want %v", res, data)
		}
	})

	os.Remove(filePath)
}

Running go test -v will execute the tests to verify everything works. The outcome should look like this if everything is working:

=== RUN   Test_addTemplate
=== RUN   Test_addTemplate/testName
New template 'testName' added!
--- PASS: Test_addTemplate (0.00s)
    --- PASS: Test_addTemplate/testName (0.00s)
=== RUN   Test_removeTemplate
=== RUN   Test_removeTemplate/testName
Template 'testName' removed!
--- PASS: Test_removeTemplate (0.00s)
    --- PASS: Test_removeTemplate/testName (0.00s)
=== RUN   Test_readFromFile
=== RUN   Test_readFromFile/test_reading
--- PASS: Test_readFromFile (0.00s)
    --- PASS: Test_readFromFile/test_reading (0.00s)
=== RUN   Test_writeToFile
=== RUN   Test_writeToFile/test_writing
--- PASS: Test_writeToFile (0.00s)
    --- PASS: Test_writeToFile/test_writing (0.00s)
PASS
ok      _/home/tyler/Code/ByteUnitsExamples/go_ssh_cli  0.004s

Utilizing

To utilize this program without having to run Go every time it can be built into a normal bash executable. Run go build tssh.go to generate an executable called tssh. Copy this generated file to the /usr/local/bin so that it's executable on your path. Now running tssh test should execute the command setup earlier.

Summarize

In this short article, we explored running a Go program to store and execute SSH commands. This is just a very brief introduction into the Go language but hopefully, it shows how useful it can be.

Tyler Moon

Published 8 months ago