Merge pull request #516 from yifan-gu/ssh_client

e2e: Add a simple ssh client and node type

Author: Yifan Gu

e2e/main_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/crypto/ssh"
+
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/pkg/api/v1"

e2e/node.go

@@ -0,0 +1,49 @@
+package e2e
+
+import (
+	"fmt"
+
+	"k8s.io/client-go/pkg/api/v1"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type Node struct {
+	apiNode *v1.Node
+}
+
+func NewNode(n *v1.Node) *Node {
+	return &Node{apiNode: n}
+}
+
+func (n *Node) ExternalIP() string {
+	var host string
+	for _, addr := range n.apiNode.Status.Addresses {
+		if addr.Type == v1.NodeExternalIP {
+			host = addr.Address
+			break
+		}
+	}
+	return host
+}
+
+func (n *Node) SSH(cmd string) (stdout, stderr []byte, err error) {
+	host := n.ExternalIP()
+	if host == "" {
+		return nil, nil, fmt.Errorf("cannot find external IP for node %q", n.apiNode.Name)
+	}
+	return sshClient.SSH(host, cmd)
+}
+
+func (n *Node) Reboot() error {
+	stdout, stderr, err := n.SSH("sudo reboot")
+	if _, ok := err.(*ssh.ExitMissingError); ok {
+		// A terminated session is perfectly normal during reboot.
+		err = nil
+	}
+
+	if err != nil {
+		return fmt.Errorf("issuing reboot command failed\nstdout:%s\nstderr:%s", stdout, stderr)
+	}
+	return nil
+}

e2e/smoke_test.go

@@ -19,7 +19,7 @@ func TestSmoke(t *testing.T) {
 	}
 	d, ok := di.(*v1beta1.Deployment)
 	if !ok {
-		t.Fatalf("expected manifest to decode into *api.Service, got %T", di)
+		t.Fatalf("expected manifest to decode into *api.deployment, got %T", di)
 	}
 	_, err = client.ExtensionsV1beta1().Deployments(namespace).Create(d)
 	if err != nil {

e2e/ssh_client.go

@@ -0,0 +1,88 @@
+package e2e
+
+import (
+	"bytes"
+	"io/ioutil"
+	"log"
+	"net"
+	"os"
+
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// The global SSH client used by the tests.
+var sshClient *SSHClient
+
+type SSHClient struct {
+	*ssh.ClientConfig
+}
+
+func InitSSHClient(keypath string) {
+	sshClient = NewSSHClientOrDie(keypath)
+}
+
+// NewSSHClientOrDie tries to create an ssh client.
+// If $SSH_AUTH_SOCK is set, the use the ssh agent to create the client,
+// otherwise read the private key directly.
+func NewSSHClientOrDie(keypath string) *SSHClient {
+	var authMethod ssh.AuthMethod
+
+	sock := os.Getenv("SSH_AUTH_SOCK")
+	if sock != "" {
+		log.Println("Creating ssh client with ssh agent")
+		sshAgent, err := net.Dial("unix", sock)
+		if err != nil {
+			panic(err)
+		}
+
+		authMethod = ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers)
+	} else {
+		log.Println("Creating ssh client with private key")
+		key, err := ioutil.ReadFile(keypath)
+		if err != nil {
+			panic(err)
+		}
+
+		signer, err := ssh.ParsePrivateKey(key)
+		if err != nil {
+			panic(err)
+		}
+
+		authMethod = ssh.PublicKeys(signer)
+	}
+
+	sshConfig := &ssh.ClientConfig{
+		User:            "core", // TODO(yifan): Assume all nodes are container linux nodes for now.
+		Auth:            []ssh.AuthMethod{authMethod},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	return &SSHClient{sshConfig}
+}
+
+func (c *SSHClient) SSH(host, cmd string) (stdout, stderr []byte, err error) {
+	client, err := ssh.Dial("tcp", host+":22", c.ClientConfig) // TODO(yifan): Assume all nodes are listening on :22 for ssh requests for now.
+	if err != nil {
+		return nil, nil, err
+	}
+	defer client.Conn.Close()
+
+	session, err := client.NewSession()
+	if err != nil {
+		return nil, nil, err
+	}
+	defer session.Close()
+
+	outBuf := bytes.NewBuffer(nil)
+	errBuf := bytes.NewBuffer(nil)
+	session.Stdout = outBuf
+	session.Stderr = errBuf
+
+	err = session.Run(cmd)
+
+	stdout = bytes.TrimSpace(outBuf.Bytes())
+	stderr = bytes.TrimSpace(errBuf.Bytes())
+
+	return stdout, stderr, err
+}

glide.lock

@@ -11,6 +11,9 @@ import:
 - package: github.com/ghodss/yaml
 - package: github.com/golang/glog
 - package: github.com/spf13/cobra
+  version: 4cdb38c072b86bf795d2c81de50784d9fdd6eb77
+- package: github.com/spf13/pflag
+  version: e57e3eeb33f795204c1ca35f56c44f83227c6e66
 - package: k8s.io/apimachinery
   version: 75b8dd260ef0469d96d578705a87cffd0e09dab8
   subpackages:
@@ -34,3 +37,10 @@ import:
   subpackages:
   - pkg/kubectl/cmd/util
   - pkg/kubectl/resource
+- package: golang.org/x/net
+  version: e90d6d0afc4c315a0d87a568ae68577cc15149a0
+- package: golang.org/x/crypto/
+  version: 0fe963104e9d1877082f8fb38f816fcd97eb1d10
+  subpackages:
+  - ssh
+  - ssh/agent