github.com/mirantis/virtlet@v1.5.2-0.20191204181327-1659b8a48e9b/tests/e2e/framework/ssh_interface.go (about)

     1  /*
     2  Copyright 2017 Mirantis
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package framework
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"io"
    23  	"strings"
    24  
    25  	"golang.org/x/crypto/ssh"
    26  
    27  	"github.com/Mirantis/virtlet/pkg/tools"
    28  )
    29  
    30  type sshInterface struct {
    31  	vmInterface *VMInterface
    32  	client      *ssh.Client
    33  	fwStopCh    chan struct{}
    34  }
    35  
    36  func newSSHInterface(vmInterface *VMInterface, user, secret string) (*sshInterface, error) {
    37  	var authMethod ssh.AuthMethod
    38  	key := trimBlock(secret)
    39  	signer, err := ssh.ParsePrivateKey([]byte(key))
    40  	if err != nil {
    41  		authMethod = ssh.Password(secret)
    42  	} else {
    43  		authMethod = ssh.PublicKeys(signer)
    44  	}
    45  
    46  	config := &ssh.ClientConfig{
    47  		User: user,
    48  		Auth: []ssh.AuthMethod{authMethod},
    49  	}
    50  
    51  	vmPod, err := vmInterface.Pod()
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  
    56  	ports := []*tools.ForwardedPort{
    57  		{RemotePort: 22},
    58  	}
    59  	fwStopCh, err := vmPod.PortForward(ports)
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  
    64  	sshClient, err := ssh.Dial("tcp", fmt.Sprintf("localhost:%d", ports[0].LocalPort), config)
    65  	if err != nil {
    66  		defer close(fwStopCh)
    67  		return nil, err
    68  	}
    69  
    70  	return &sshInterface{
    71  		vmInterface: vmInterface,
    72  		client:      sshClient,
    73  		fwStopCh:    fwStopCh,
    74  	}, nil
    75  }
    76  
    77  func (si *sshInterface) Run(stdin io.Reader, stdout, stderr io.Writer, command ...string) error {
    78  	cmd, err := si.Start(stdin, stdout, stderr, command...)
    79  	if err != nil {
    80  		return err
    81  	}
    82  
    83  	return cmd.Wait()
    84  }
    85  
    86  func (si *sshInterface) Start(stdin io.Reader, stdout, stderr io.Writer, command ...string) (Command, error) {
    87  	session, err := si.client.NewSession()
    88  	if err != nil {
    89  		return nil, err
    90  	}
    91  
    92  	session.Stdout = stdout
    93  	session.Stderr = stderr
    94  	session.Stdin = stdin
    95  
    96  	for i, arg := range command {
    97  		if i == 0 {
    98  			continue
    99  		}
   100  		command[i] = fmt.Sprintf("'%s'", strings.Replace(arg, "'", "\\'", -1))
   101  	}
   102  	if err := session.Start(strings.Join(command, " ")); err != nil {
   103  		defer session.Close()
   104  		return nil, err
   105  	}
   106  
   107  	return sshCommand{session: session}, err
   108  }
   109  
   110  func (si *sshInterface) Close() error {
   111  	if si.client != nil {
   112  		defer close(si.fwStopCh)
   113  		err := si.client.Close()
   114  		si.client = nil
   115  		return err
   116  	}
   117  	return nil
   118  }
   119  
   120  // Logs is a placeholder for fulfilling Executor interface
   121  func (*sshInterface) Logs() (string, error) {
   122  	return "", errors.New("not implemented")
   123  }
   124  
   125  type sshCommand struct {
   126  	session *ssh.Session
   127  }
   128  
   129  var _ Command = &sshCommand{}
   130  
   131  func (sc sshCommand) Wait() error {
   132  	err := sc.session.Wait()
   133  	defer sc.session.Close()
   134  	if err != nil {
   135  		if s, ok := err.(*ssh.ExitError); ok {
   136  			return CommandError{ExitCode: s.ExitStatus()}
   137  		}
   138  		return err
   139  	}
   140  
   141  	return nil
   142  }
   143  
   144  func (sc sshCommand) Kill() error {
   145  	defer sc.session.Close()
   146  	return sc.session.Signal(ssh.SIGKILL)
   147  }