github.com/cloudfoundry-attic/ltc@v0.0.0-20151123212628-098adc7919fc/ssh/command_factory/command_factory.go

package command_factory

import (
	"fmt"
	"runtime"
	"strings"

	"github.com/cloudfoundry-incubator/ltc/app_examiner"
	config_package "github.com/cloudfoundry-incubator/ltc/config"
	"github.com/cloudfoundry-incubator/ltc/exit_handler"
	"github.com/cloudfoundry-incubator/ltc/exit_handler/exit_codes"
	"github.com/cloudfoundry-incubator/ltc/terminal"
	"github.com/codegangsta/cli"
)

type SSHCommandFactory struct {
	config      *config_package.Config
	ui          terminal.UI
	exitHandler exit_handler.ExitHandler
	appExaminer app_examiner.AppExaminer

	secureShell SSH
}

//go:generate counterfeiter -o mocks/fake_ssh.go . SSH
type SSH interface {
	Connect(appName string, instanceIndex int, config *config_package.Config) error
	Forward(localAddress, remoteAddress string) error
	Shell(command string, ptyDespired bool) error
}

func NewSSHCommandFactory(config *config_package.Config, ui terminal.UI, exitHandler exit_handler.ExitHandler, appExaminer app_examiner.AppExaminer, secureShell SSH) *SSHCommandFactory {
	return &SSHCommandFactory{config, ui, exitHandler, appExaminer, secureShell}
}

func (f *SSHCommandFactory) MakeSSHCommand() cli.Command {
	helpText := "ltc ssh <app-name> [[--] optional command with args]\n\n   If a command is specified, no interactive shell will be provided.\n   A \"--\" token should be provided to avoid parsing of command flags.\n"
	if runtime.GOOS == "windows" {
		helpText = helpText + "\n   ltc ssh requires a pseudo-terminal and will not function properly from the default Windows command shells (Cmd, Powershell). We recommend using ltc ssh non-interactively: ltc ssh -- ps auxw\n"
	}

	return cli.Command{
		Name:        "ssh",
		Usage:       "Connects to a running app",
		Description: helpText,
		Action:      f.ssh,
		Flags: []cli.Flag{
			cli.IntFlag{
				Name:  "index, i",
				Usage: "Connects to specified instance index",
				Value: 0,
			},
			cli.StringFlag{
				Name:  "L",
				Usage: "Listens on specified local address/port and forwards connections to specified remote address/port\n     \te.g. ltc ssh <app-name> -L [localhost:]<local-port>:<remote-host>:<remote-port>",
			},
			cli.BoolFlag{
				Name:  "N",
				Usage: "Disables the interactive shell when forwarding connections with -L",
			},
			cli.BoolFlag{
				Name:  "T",
				Usage: "Disables pseudo-tty allocation",
			},
			cli.BoolFlag{
				Name:  "t",
				Usage: "Enables pseudo-tty allocation",
			},
		},
	}
}

func (f *SSHCommandFactory) ssh(context *cli.Context) {
	instanceIndex := context.Int("index")
	localForward := context.String("L")
	noShell := context.Bool("N")
	forceNoPTY := context.Bool("T")

	appName := context.Args().First()

	if appName == "" {
		f.ui.SayIncorrectUsage("Please input a valid <app-name>")
		f.exitHandler.Exit(exit_codes.InvalidSyntax)
		return
	}

	appInfo, err := f.appExaminer.AppStatus(appName)
	if err != nil {
		f.ui.SayLine("App " + appName + " not found.")
		f.exitHandler.Exit(exit_codes.CommandFailed)
		return
	}
	if instanceIndex < 0 || instanceIndex >= appInfo.ActualRunningInstances {
		f.ui.SayLine("Instance %s/%d does not exist.", appName, instanceIndex)
		f.exitHandler.Exit(exit_codes.CommandFailed)
		return
	}

	if err := f.secureShell.Connect(appName, instanceIndex, f.config); err != nil {
		f.ui.SayLine("Error connecting to %s/%d: %s", appName, instanceIndex, err)
		f.exitHandler.Exit(exit_codes.CommandFailed)
		return
	}

	command := ""
	if len(context.Args()) > 1 {
		start := 1
		if context.Args().Get(1) == "--" {
			start = 2
		}
		command = strings.Join(context.Args()[start:len(context.Args())], " ")
	}

	if localForward != "" && noShell {
		f.forward(localForward, appName, instanceIndex)
	} else if localForward != "" {
		go f.forward(localForward, appName, instanceIndex)
		f.shell(command, !forceNoPTY, appName, instanceIndex)
	} else {
		f.shell(command, !forceNoPTY, appName, instanceIndex)
	}
}

func (f *SSHCommandFactory) forward(localForward, appName string, instanceIndex int) {
	var localHost, localPort, remoteHost, remotePort string

	parts := strings.Split(localForward, ":")

	switch len(parts) {
	case 3:
		localHost, localPort, remoteHost, remotePort = "localhost", parts[0], parts[1], parts[2]
	case 4:
		localHost, localPort, remoteHost, remotePort = parts[0], parts[1], parts[2], parts[3]
	default:
		f.ui.SayIncorrectUsage("-L expects [localhost:]<local-port>:<remote-host>:<remote-port>")
		f.exitHandler.Exit(exit_codes.InvalidSyntax)
		return
	}

	if localPort == "0" {
		f.ui.SayIncorrectUsage("-L expects [localhost:]<local-port>:<remote-host>:<remote-port>")
		f.exitHandler.Exit(exit_codes.InvalidSyntax)
		return
	}

	localAddr := fmt.Sprintf("%s:%s", localHost, localPort)
	remoteAddr := fmt.Sprintf("%s:%s", remoteHost, remotePort)

	f.ui.SayLine("Forwarding %s to %s via %s/%d at %s", localAddr, remoteAddr, appName, instanceIndex, f.config.Target())

	if err := f.secureShell.Forward(localAddr, remoteAddr); err != nil {
		f.ui.SayLine("Error connecting to %s/%d: %s", appName, instanceIndex, err)
		f.exitHandler.Exit(exit_codes.CommandFailed)
		return
	}
}

func (f *SSHCommandFactory) shell(command string, ptyDesired bool, appName string, instanceIndex int) {
	if command == "" {
		f.ui.SayLine("Connecting to %s/%d at %s", appName, instanceIndex, f.config.Target())
	}

	if err := f.secureShell.Shell(command, ptyDesired); err != nil {
		f.ui.SayLine("Error connecting to %s/%d: %s", appName, instanceIndex, err)
		f.exitHandler.Exit(exit_codes.CommandFailed)
		return
	}
}