cmd/testutil/run.go

// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"io/ioutil"
	"os"
	"time"

	"github.com/google/subcommands"

	"fuchsia.googlesource.com/testing/netboot"
	"fuchsia.googlesource.com/testing/ssh"
)

type runCmd struct {
	key string
	host string
	kernel string
	ramdisk string
	user string
	timeout time.Duration
}

func (*runCmd) Name() string { return "run" }
func (*runCmd) Synopsis() string { return "Run a command." }
func (*runCmd) Usage() string {
	return `run <command>:
  Run a command on the target device.
`
}

func (b *runCmd) SetFlags(f *flag.FlagSet) {
	f.StringVar(&b.key, "key", "", "path to the key")
	f.StringVar(&b.host, "host", "", "target hostname")
	f.StringVar(&b.kernel, "kernel", "", "kernel image")
	f.StringVar(&b.ramdisk, "ramdisk", "", "ramdisk image")
	f.StringVar(&b.user, "user", "fuchsia", "user identity")
	f.DurationVar(&b.timeout, "timeout", 1*time.Minute, "timeout")
}

func (b *runCmd) Execute(ctx context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	if len(f.Args()) == 0 {
		return subcommands.ExitUsageError
	}

	if err := b.run(ctx, f.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "error: %v", err)
		return subcommands.ExitFailure
	}

	return subcommands.ExitSuccess
}

func (b *runCmd) run(ctx context.Context, args []string) error {
	key, err := ioutil.ReadFile(b.key)
	if err != nil {
		return fmt.Errorf("cannot read key: %v", err)
	}

	n := &netboot.Netboot{Timeout: b.timeout}
	var addr *net.UDPAddr

	if b.host != "" {
		addr, err = n.Discover(ctx, b.host, true)
		if err != nil {
			return fmt.Errorf("cannot find device `%s`: %v\n", b.host, err)
		}
	} else {
		addr, err = n.Beacon(ctx, )
		if err != nil {
			return fmt.Errorf("cannot get beacon: %v\n", err)
		}
	}

	addr.Port = ssh.Port

	c, err := ssh.NewConnection(fmt.Sprintf("%v", addr), b.user, key)
	if err != nil {
		return fmt.Errorf("cannot connect to target: %v", err)
	}

	ctx, cancel := context.WithTimeout(ctx, b.timeout)
	defer cancel()

	err = c.Run(ctx, args[0], args[1:]...)
	if err != nil {
		return fmt.Errorf("failed to execute command: %v", err)
	}

	return nil
}