NetOpsCoding_Advent_Calendar_2016_sample.go

package main

import (
    "bytes"
    "fmt"
    "io"
    "os"

    "golang.org/x/crypto/ssh"
)

const (
    addr     = "192.0.2.1:22"
    username = "user"
    password = "password"
)

type KeyboardInteractive map[string]string

func (ki KeyboardInteractive) Challenge(user, instruction string, questions []string, echos []bool) ([]string, error) {
    var answers []string

    for _, q := range questions {
        answers = append(answers, ki[q])
    }

    return answers, nil
}

type SSHConn struct {
    session *ssh.Session
    stdin   io.WriteCloser
    stdout  io.Reader
    stderr  io.Reader
}

func (conn *SSHConn) Send(msg string) (n int, err error) {
    return conn.stdin.Write([]byte(msg + "\n"))
}

func (conn *SSHConn) Recv(suff string) (string, error) {
    var result bytes.Buffer
    buff := make([]byte, 1024)
    for {
        n, err := conn.stdout.Read(buff)
        if err != io.EOF && err != nil {
            return "", err
        }

        result.Write(buff[:n])
        if err == io.EOF || bytes.HasSuffix(buff[:n], []byte(suff)) {
            break
        }
    }

    return result.String(), nil
}

func (conn *SSHConn) Close() error {
    return conn.session.Close()
}

func setupSSHConn(addr, username, password string) (*SSHConn, error) {
    client, err := _setupSSHClient(addr, username, password)
    if err != nil {
        return nil, err
    }

    return _setupSSHSession(client)
}

func _setupSSHClient(addr, username, password string) (*ssh.Client, error) {
    config := ssh.Config{}
    config.SetDefaults()
    config.Ciphers = append(config.Ciphers, "aes128-cbc")

    answers := KeyboardInteractive(map[string]string{
        "Password: ": password,
    })

    clientConfig := &ssh.ClientConfig{
        Config: config,
        User:   username,
        Auth: []ssh.AuthMethod{
            ssh.Password(password),
            ssh.KeyboardInteractive(
                answers.Challenge,
            ),
        },
    }

    client, err := ssh.Dial("tcp", addr, clientConfig)
    if err != nil {
        return nil, err
    }

    return client, nil
}

func _setupSSHSession(client *ssh.Client) (*SSHConn, error) {
    session, err := client.NewSession()
    if err != nil {
        return nil, err
    }

    stdin, err := session.StdinPipe()
    if err != nil {
        return nil, err
    }

    stdout, err := session.StdoutPipe()
    if err != nil {
        return nil, err
    }

    stderr, err := session.StderrPipe()
    if err != nil {
        return nil, err
    }

    modes := ssh.TerminalModes{
        ssh.ECHO:          0,
        ssh.ECHOCTL:       0,
        ssh.TTY_OP_ISPEED: 115200,
        ssh.TTY_OP_OSPEED: 115200,
    }

    err = session.RequestPty("xterm", 80, 80, modes)
    if err != nil {
        return nil, err
    }

    err = session.Shell()
    if err != nil {
        return nil, err
    }

    return &SSHConn{
        session: session,
        stdin:   stdin,
        stdout:  stdout,
        stderr:  stderr,
    }, nil
}

func main() {
    sshConn, err := setupSSHConn(addr, username, password)
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }
    defer sshConn.Close()

    _, err = sshConn.Send("set cli screen-length 0")
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }

    result, err := sshConn.Recv("> ")
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }
    fmt.Print(result)

    _, err = sshConn.Send("show sys uptime")
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }

    result, err = sshConn.Recv("> ")
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }
    fmt.Print(result)

    _, err = sshConn.Send("exit")
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }

    result, err = sshConn.Recv("> ")
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }
    fmt.Print(result)
}