sigs.k8s.io/cluster-api-provider-aws@v1.5.5/test/e2e/shared/exec.go (about)

     1  //go:build e2e
     2  // +build e2e
     3  
     4  /*
     5  Copyright 2020 The Kubernetes Authors.
     6  
     7  Licensed under the Apache License, Version 2.0 (the "License");
     8  you may not use this file except in compliance with the License.
     9  You may obtain a copy of the License at
    10  
    11  	http://www.apache.org/licenses/LICENSE-2.0
    12  
    13  Unless required by applicable law or agreed to in writing, software
    14  distributed under the License is distributed on an "AS IS" BASIS,
    15  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    16  See the License for the specific language governing permissions and
    17  limitations under the License.
    18  */
    19  
    20  package shared
    21  
    22  import (
    23  	"context"
    24  	"encoding/json"
    25  	"fmt"
    26  	"os"
    27  	"path"
    28  	"path/filepath"
    29  	"regexp"
    30  	"time"
    31  
    32  	"github.com/aws/aws-sdk-go/aws"
    33  	"github.com/aws/aws-sdk-go/service/ec2"
    34  	"github.com/aws/aws-sdk-go/service/ssm"
    35  	expect "github.com/google/goexpect"
    36  )
    37  
    38  type instance struct {
    39  	name       string
    40  	instanceID string
    41  }
    42  
    43  // allMachines gets all EC2 instances at once, to save on DescribeInstances calls.
    44  func allMachines(ctx context.Context, e2eCtx *E2EContext) ([]instance, error) {
    45  	ec2Svc := ec2.New(e2eCtx.AWSSession)
    46  	resp, err := ec2Svc.DescribeInstancesWithContext(ctx, &ec2.DescribeInstancesInput{})
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  	if len(resp.Reservations) == 0 || len(resp.Reservations[0].Instances) == 0 {
    51  		return nil, fmt.Errorf("no machines found")
    52  	}
    53  	instances := []instance{}
    54  	for _, r := range resp.Reservations {
    55  		for _, i := range r.Instances {
    56  			tags := i.Tags
    57  			name := ""
    58  			for _, t := range tags {
    59  				if aws.StringValue(t.Key) == "Name" {
    60  					name = aws.StringValue(t.Value)
    61  				}
    62  			}
    63  			if name == "" {
    64  				continue
    65  			}
    66  			instances = append(instances,
    67  				instance{
    68  					name:       name,
    69  					instanceID: aws.StringValue(i.InstanceId),
    70  				},
    71  			)
    72  		}
    73  	}
    74  	return instances, nil
    75  }
    76  
    77  type command struct {
    78  	title string
    79  	cmd   string
    80  }
    81  
    82  // commandsForMachine opens a terminal connection using AWS SSM Session Manager
    83  // and executes the given commands, outputting the results to a file for each.
    84  func commandsForMachine(ctx context.Context, e2eCtx *E2EContext, f *os.File, instanceID string, commands []command) {
    85  	ssmSvc := ssm.New(e2eCtx.BootstrapUserAWSSession)
    86  	sess, err := ssmSvc.StartSessionWithContext(ctx, &ssm.StartSessionInput{
    87  		Target: aws.String(instanceID),
    88  	})
    89  	if err != nil {
    90  		fmt.Fprintf(f, "unable to start session: err=%s", err)
    91  		return
    92  	}
    93  	defer func() {
    94  		if _, err := ssmSvc.TerminateSession(&ssm.TerminateSessionInput{SessionId: sess.SessionId}); err != nil {
    95  			fmt.Fprintf(f, "unable to terminate session: err=%s", err)
    96  		}
    97  	}()
    98  	sessionToken, err := json.Marshal(sess)
    99  	if err != nil {
   100  		fmt.Fprintf(f, "unable to marshal session: err=%s", err)
   101  		return
   102  	}
   103  	cmdLine := fmt.Sprintf("session-manager-plugin %s %s StartSession %s", string(sessionToken), *ssmSvc.Client.Config.Region, ssmSvc.Client.Endpoint)
   104  	e, _, err := expect.Spawn(cmdLine, -1)
   105  	if err != nil {
   106  		fmt.Fprintf(f, "unable to spawn AWS SSM Session Manager plugin: %s", err)
   107  		return
   108  	}
   109  	defer e.Close()
   110  	shellStart := regexp.MustCompile(`\n\$`)
   111  	if result, _, err := e.Expect(shellStart, 10*time.Second); err != nil {
   112  		fmt.Fprintf(f, "did not find shell: err=%s, output=%s", err, result)
   113  		return
   114  	}
   115  	for _, c := range commands {
   116  		logFile := path.Join(filepath.Dir(f.Name()), c.title+".log")
   117  		if err := e.Send("sudo " + c.cmd + "\n"); err != nil {
   118  			fmt.Fprintf(f, "unable to send command: err=%s", err)
   119  			return
   120  		}
   121  		result, _, err := e.Expect(shellStart, 20*time.Second)
   122  		if err := os.WriteFile(logFile, []byte(result), os.ModePerm); err != nil {
   123  			fmt.Fprintf(f, "error writing log file: err=%s", err)
   124  			return
   125  		}
   126  		if err != nil {
   127  			fmt.Fprintf(f, "error awaiting command output: err=%s", err)
   128  			return
   129  		}
   130  	}
   131  }