3
\$\begingroup\$

I'm trying to learn a bit about Golang and decided to try out building an SSH address book command-line app. I used a YAML library to serialize and store data in a configuration file. The app has a few command-line options to view, delete, make, and run SSH address entries. It also runs the ssh command given an entry name and no arguments.

I'm looking for some criticism on things I should improve.

package main

import (
    "flag"
    "fmt"
    "golang.org/x/crypto/ssh"
    "golang.org/x/crypto/ssh/terminal"
    "gopkg.in/yaml.v2"
    "io"
    "io/ioutil"
    "log"
    "os"
    "os/exec"
)

// Sturct that contains the list of all SSHEntry structs
type EntriesList struct {
    SSHEntries []SSHEntry `yaml:"sshentries"`
}

// Struct that holds the name of the SSHCmd entry and the SSHCmd struct itself
type SSHEntry struct {
    Name    string `yaml:"name"`
    Command SSHCmd `yaml:"command"`
}

// Struct that contains the Address, Port number, and Username for executing ssh commands
type SSHCmd struct {
    Addr  string `yaml:"addr"`
    Port  int    `yaml:"port"`
    Uname string `yaml:"uname"`
}

func main() {
    MAKE := flag.Bool("make", false, "Specify whether you want to make a new entry")
    CONF := flag.String("conf", "/misc/sshbook-g/sshbook.yml", "Specify location of configuration file")
    VIEW := flag.Bool("view", false, "View previous entries")
    CLEAN := flag.Bool("clean", false, "Delete all stored entries")
    CHECK := flag.String("check", "", "Check for entry stored in yaml file by entry name")
    DELETE := flag.Bool("delete", false, "Delete an entry")

    flag.Parse()
    if flag.NFlag() > 0 {
        if *MAKE {
            entries := readConf(*CONF)
            addEntry(entries)
            writeEntries(entries, *CONF)
        } else if *VIEW {
            viewEntries(*CONF)
        } else if *CLEAN {
            cleanEntries(*CONF)
        } else if len(*CHECK) > 0 {
            entries := readConf(*CONF)
            if searchEntriesList(entries, *CHECK) != nil {
                fmt.Printf("%s is available\n", *CHECK)
            } else {
                fmt.Printf("%s is not available\n", *CHECK)
            }
        } else if *DELETE {
            entries := readConf(*CONF)
            deleteAnEntry(entries, *CONF)
        }
        return
    }

    if len(os.Args) < 1 {
        panic("Need an argument")
    }
    entries := readConf(*CONF)
    entry_name := os.Args[1]
    runnable_entry := searchEntriesList(entries, entry_name)
    if runnable_entry != nil {
        run(runnable_entry.Command)
    } else {
        fmt.Println("That entry is not available")
    }
}

// Allows a user to delete an entry, First prints out the whole list 
// and has the user to specify the name of the entry to be deleted
func deleteAnEntry(entries *EntriesList, CONF string) {
    viewEntries(CONF)
    fmt.Println("Enter the name of the entry you wish to delete:")
    var answer string
    fmt.Scanf("%s", &answer)
    toDelete := searchEntriesListIndex(entries, answer)
    if toDelete == -1 {
        fmt.Println("Cannot find that entry")
        return
    }
    entries.SSHEntries = append(entries.SSHEntries[:toDelete], entries.SSHEntries[toDelete+1:]...)
    writeEntries(entries, CONF)
}

// Calls makeEntry to create a new SSHEntry struct and adds it to the master list of current entries.
func addEntry(entries *EntriesList) {
    newEntry := makeEntry()
    entries.SSHEntries = append(entries.SSHEntries, *newEntry)
}

// Creates a new struct via prompting user for imformation about ssh command they wish to enter
// Returns pointer to new SSHEntry struct
func makeEntry() (entry *SSHEntry) {
    var uname, addr, entryName string
    var port int
    port = -1
    fmt.Println("Enter the data for a new entry")
    fmt.Println("Enter the username: ")
    fmt.Scanf("%s\n", &uname)
    fmt.Println("Enter the port number: ")
    fmt.Println("(Leave empty for default port)")
    fmt.Scanf("%d\n", &port)
    fmt.Println("Enter the address: ")
    fmt.Scanf("%s\n", &addr)
    fmt.Println("Enter the entry name you would like to use: ")
    fmt.Scanf("%s\n", &entryName)
    // If port is unchanged, then the user picked the default option
    if port < 0 {
        port = 22
    }
    return &SSHEntry{entryName, SSHCmd{addr, port, uname}}
}

// Generic error checking function
func check(e error) {
    if e != nil && e != io.EOF {
        log.Fatal(e)
    }
}

// Serializes entries list and writes it to the configuration file using ioutil's WriteFile function
func writeEntries(entries *EntriesList, CONF string) {
    bodyBytes, err := yaml.Marshal(entries)
    check(err)
    ioutil.WriteFile(CONF, bodyBytes, 0644)
    check(err)
}

// Deletes entry list by writing an empyt byte[] to the configuration file
func cleanEntries(CONF string) {
    var answer string
    fmt.Println("This will clean out all the currently stored entries. Are you sure? y/n")
    fmt.Scanf("%s\n", &answer)
    if answer != "y" {
        return
    }
    ioutil.WriteFile(CONF, []byte(""), 0644)
}

// Prints entries as yaml data by running the shell command `cat` on the configuratoin file and printing it's output.
func viewEntries(CONF string) {
    catCmd := exec.Command("cat", CONF)
    output, err := catCmd.Output()
    check(err)
    fmt.Println(string(output))
}

// Performs linear search on entries list for a given name.
// Returns nil if no entry names match searchString.
func searchEntriesList(entries *EntriesList, searchString string) *SSHEntry {
    for i := 0; i < len(entries.SSHEntries); i += 1 {
        if entries.SSHEntries[i].Name == searchString {
            return &entries.SSHEntries[i]
        }
    }
    return nil
}

// Identical to searchEntriesList but returns index in entries list rather than struct pointer.
// returns -1 if no results are found.
func searchEntriesListIndex(entries *EntriesList, searchString string) int {
    for i := 0; i < len(entries.SSHEntries); i += 1 {
        if entries.SSHEntries[i].Name == searchString {
            return i
        }
    }
    return -1
}

// Reads configuration file and deserializes yaml data into native structs.
// Returns pointer to list of all entries.
func readConf(CONF string) *EntriesList {
    confiFile, err := ioutil.ReadFile(CONF)
    check(err)
    var entries EntriesList
    err = yaml.Unmarshal(confiFile, &entries)
    check(err)
    return &entries
}

// Runs the ssh command using the given SSHCmd struct to gather necessary data.
// Currently only supports using passwords as authentication methods.
// Waits for session to end then returns
func run(command SSHCmd) error {
    password := askForPassword()
    config := ssh.ClientConfig{
        User: command.Uname,
        Auth: []ssh.AuthMethod{ssh.Password(password)},
    }
    conn, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", command.Addr, command.Port), &config)
    check(err)
    session, err := conn.NewSession()
    check(err)
    defer session.Close()
    session.Stdout = os.Stdout
    session.Stderr = os.Stderr
    session.Stdin = os.Stdin

    modes := ssh.TerminalModes{
        ssh.ECHO:          0,
        ssh.ECHOCTL:       0,
        ssh.TTY_OP_ISPEED: 14400,
        ssh.TTY_OP_OSPEED: 14400,
    }
    if err = session.RequestPty("xterm", 80, 40, modes); err != nil {
        check(err)
    }
    if err = session.Shell(); err != nil {
        check(err)
    }

    return session.Wait()
}

// Asks for password
// Returns password as string
func askForPassword() string {
    fmt.Println("Enter your password:")
    pwbytes, err := terminal.ReadPassword(0)
    check(err)
    return string(pwbytes)
}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.