// setProxyCommand sets the proxy command option.
func (c *SSHCommon) setProxyCommand(options *ssh.Options) error {
	apiServerHost, _, err := net.SplitHostPort(c.apiAddr)
	if err != nil {
		return errors.Errorf("failed to get proxy address: %v", err)
	}
	juju, err := getJujuExecutable()
	if err != nil {
		return errors.Errorf("failed to get juju executable path: %v", err)
	}

	// TODO(mjs) 2016-05-09 LP #1579592 - It would be good to check the
	// host key of the controller machine being used for proxying
	// here. This isn't too serious as all traffic passing through the
	// controller host is encrypted and the host key of the ultimate
	// target host is verified but it would still be better to perform
	// this extra level of checking.
	options.SetProxyCommand(
		juju, "ssh",
		"--proxy=false",
		"--no-host-key-checks",
		"--pty=false",
		"[email protected]"+apiServerHost,
		"-q",
		"nc %h %p",
	)
	return nil
}

func (s *SSHGoCryptoCommandSuite) TestProxyCommand(c *gc.C) {
	realNetcat, err := exec.LookPath("nc")
	if err != nil {
		c.Skip("skipping test, couldn't find netcat: %v")
		return
	}
	netcat := filepath.Join(c.MkDir(), "nc")
	err = ioutil.WriteFile(netcat, []byte("#!/bin/sh\necho $0 \"[email protected]\" > $0.args && exec "+realNetcat+" \"[email protected]\""), 0755)
	c.Assert(err, jc.ErrorIsNil)

	private, _, err := ssh.GenerateKey("test-server")
	c.Assert(err, jc.ErrorIsNil)
	key, err := cryptossh.ParsePrivateKey([]byte(private))
	client, err := ssh.NewGoCryptoClient(key)
	c.Assert(err, jc.ErrorIsNil)
	server := newServer(c)
	var opts ssh.Options
	port := server.listener.Addr().(*net.TCPAddr).Port
	opts.SetProxyCommand(netcat, "-q0", "%h", "%p")
	opts.SetPort(port)
	cmd := client.Command("127.0.0.1", testCommand, &opts)
	server.cfg.PublicKeyCallback = func(_ cryptossh.ConnMetadata, pubkey cryptossh.PublicKey) (*cryptossh.Permissions, error) {
		return nil, nil
	}
	go server.run(c)
	out, err := cmd.Output()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(string(out), gc.Equals, "abc value\n")
	// Ensure the proxy command was executed with the appropriate arguments.
	data, err := ioutil.ReadFile(netcat + ".args")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(string(data), gc.Equals, fmt.Sprintf("%s -q0 127.0.0.1 %v\n", netcat, port))
}

func (s *SSHCommandSuite) TestSetStrictHostKeyChecking(c *gc.C) {
	commandPattern := fmt.Sprintf("%s%%s -o PasswordAuthentication no -o ServerAliveInterval 30 localhost %s 123",
		s.fakessh, echoCommand)

	tests := []struct {
		input    ssh.StrictHostChecksOption
		expected string
	}{
		{ssh.StrictHostChecksNo, "no"},
		{ssh.StrictHostChecksYes, "yes"},
		{ssh.StrictHostChecksAsk, "ask"},
		{ssh.StrictHostChecksUnset, ""},
		{ssh.StrictHostChecksOption(999), ""},
	}
	for _, t := range tests {
		var opts ssh.Options
		opts.SetStrictHostKeyChecking(t.input)
		expectedOpt := ""
		if t.expected != "" {
			expectedOpt = " -o StrictHostKeyChecking " + t.expected
		}
		s.assertCommandArgs(c, s.commandOptions([]string{echoCommand, "123"}, &opts),
			fmt.Sprintf(commandPattern, expectedOpt))
	}
}

func (s *SSHCommandSuite) TestCommandPort(c *gc.C) {
	var opts ssh.Options
	opts.SetPort(2022)
	s.assertCommandArgs(c, s.commandOptions([]string{echoCommand, "123"}, &opts),
		fmt.Sprintf("%s -o StrictHostKeyChecking no -o PasswordAuthentication no -o ServerAliveInterval 30 -p 2022 localhost %s 123",
			s.fakessh, echoCommand),
	)
}

func (s *SSHCommandSuite) TestCommandSetKnownHostsFile(c *gc.C) {
	var opts ssh.Options
	opts.SetKnownHostsFile("/tmp/known hosts")
	s.assertCommandArgs(c, s.commandOptions([]string{echoCommand, "123"}, &opts),
		fmt.Sprintf("%s -o StrictHostKeyChecking no -o PasswordAuthentication no -o ServerAliveInterval 30 -o UserKnownHostsFile \"/tmp/known hosts\" localhost %s 123",
			s.fakessh, echoCommand),
	)
}

// NewSshInstanceConfigurator creates new sshInstanceConfigurator.
func NewSshInstanceConfigurator(host string) InstanceConfigurator {
	options := ssh.Options{}
	options.SetIdentities("/var/lib/juju/system-identity")
	return &sshInstanceConfigurator{
		client:  ssh.DefaultClient,
		host:    "[email protected]" + host,
		options: &options,
	}
}

func (s *SSHCommandSuite) TestCopy(c *gc.C) {
	var opts ssh.Options
	opts.EnablePTY()
	opts.AllowPasswordAuthentication()
	opts.SetIdentities("x", "y")
	opts.SetPort(2022)
	err := s.client.Copy([]string{"/tmp/blah", "[email protected]:baz"}, &opts)
	c.Assert(err, jc.ErrorIsNil)
	out, err := ioutil.ReadFile(s.fakescp + ".args")
	c.Assert(err, jc.ErrorIsNil)
	// EnablePTY has no effect for Copy
	c.Assert(string(out), gc.Equals, s.fakescp+" -o StrictHostKeyChecking no -o ServerAliveInterval 30 -i x -i y -P 2022 /tmp/blah [email protected]:baz\n")

	// Try passing extra args
	err = s.client.Copy([]string{"/tmp/blah", "[email protected]:baz", "-r", "-v"}, &opts)
	c.Assert(err, jc.ErrorIsNil)
	out, err = ioutil.ReadFile(s.fakescp + ".args")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(string(out), gc.Equals, s.fakescp+" -o StrictHostKeyChecking no -o ServerAliveInterval 30 -i x -i y -P 2022 /tmp/blah [email protected]:baz -r -v\n")

	// Try interspersing extra args
	err = s.client.Copy([]string{"-r", "/tmp/blah", "-v", "[email protected]:baz"}, &opts)
	c.Assert(err, jc.ErrorIsNil)
	out, err = ioutil.ReadFile(s.fakescp + ".args")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(string(out), gc.Equals, s.fakescp+" -o StrictHostKeyChecking no -o ServerAliveInterval 30 -i x -i y -P 2022 -r /tmp/blah -v [email protected]:baz\n")
}

// setProxyCommand sets the proxy command option.
func (c *SSHCommon) setProxyCommand(options *ssh.Options) error {
	apiServerHost, _, err := net.SplitHostPort(c.apiAddr)
	if err != nil {
		return fmt.Errorf("failed to get proxy address: %v", err)
	}
	juju, err := getJujuExecutable()
	if err != nil {
		return fmt.Errorf("failed to get juju executable path: %v", err)
	}
	options.SetProxyCommand(juju, "ssh", "--proxy=false", "--pty=false", apiServerHost, "nc", "%h", "%p")
	return nil
}

// runViaSSH runs script in the remote machine with address addr.
func runViaSSH(addr string, script string) error {
	// This is taken from cmd/juju/ssh.go there is no other clear way to set user
	userAddr := "[email protected]" + addr
	sshOptions := ssh.Options{}
	sshOptions.SetIdentities("/var/lib/juju/system-identity")
	userCmd := sshCommand(userAddr, []string{"sudo", "-n", "bash", "-c " + utils.ShQuote(script)}, &sshOptions)
	var stderrBuf bytes.Buffer
	userCmd.Stderr = &stderrBuf
	if err := userCmd.Run(); err != nil {
		return errors.Annotatef(err, "ssh command failed: %q", stderrBuf.String())
	}
	return nil
}

func (s *SSHCommandSuite) TestCommandClientKeys(c *gc.C) {
	defer overrideGenerateKey(c).Restore()
	clientKeysDir := c.MkDir()
	defer ssh.ClearClientKeys()
	err := ssh.LoadClientKeys(clientKeysDir)
	c.Assert(err, jc.ErrorIsNil)
	ck := filepath.Join(clientKeysDir, "juju_id_rsa")
	var opts ssh.Options
	opts.SetIdentities("x", "y")
	s.assertCommandArgs(c, s.commandOptions([]string{echoCommand, "123"}, &opts),
		fmt.Sprintf("%s -o StrictHostKeyChecking no -o PasswordAuthentication no -o ServerAliveInterval 30 -i x -i y -i %s localhost %s 123",
			s.fakessh, ck, echoCommand),
	)
}

// getSSHOptions configures and returns SSH options and proxy settings.
func (c *SSHCommon) getSSHOptions(enablePty bool) (*ssh.Options, error) {
	var options ssh.Options

	// TODO(waigani) do not save fingerprint only until this bug is addressed:
	// lp:892552. Also see lp:1334481.
	options.SetKnownHostsFile("/dev/null")
	if enablePty {
		options.EnablePTY()
	}
	var err error
	if c.proxy, err = c.proxySSH(); err != nil {
		return nil, err
	} else if c.proxy {
		if err := c.setProxyCommand(&options); err != nil {
			return nil, err
		}
	}
	return &options, nil
}

func (s *SSHCommandSuite) TestCommandDefaultIdentities(c *gc.C) {
	var opts ssh.Options
	tempdir := c.MkDir()
	def1 := filepath.Join(tempdir, "def1")
	def2 := filepath.Join(tempdir, "def2")
	s.PatchValue(ssh.DefaultIdentities, []string{def1, def2})
	// If no identities are specified, then the defaults aren't added.
	s.assertCommandArgs(c, s.commandOptions([]string{echoCommand, "123"}, &opts),
		fmt.Sprintf("%s -o StrictHostKeyChecking no -o PasswordAuthentication no -o ServerAliveInterval 30 localhost %s 123",
			s.fakessh, echoCommand),
	)
	// If identities are specified, then the defaults are must added.
	// Only the defaults that exist on disk will be added.
	err := ioutil.WriteFile(def2, nil, 0644)
	c.Assert(err, jc.ErrorIsNil)
	opts.SetIdentities("x", "y")
	s.assertCommandArgs(c, s.commandOptions([]string{echoCommand, "123"}, &opts),
		fmt.Sprintf("%s -o StrictHostKeyChecking no -o PasswordAuthentication no -o ServerAliveInterval 30 -i x -i y -i %s localhost %s 123",
			s.fakessh, def2, echoCommand),
	)
}

func (s *SSHGoCryptoCommandSuite) TestCommand(c *gc.C) {
	private, _, err := ssh.GenerateKey("test-server")
	c.Assert(err, jc.ErrorIsNil)
	key, err := cryptossh.ParsePrivateKey([]byte(private))
	client, err := ssh.NewGoCryptoClient(key)
	c.Assert(err, jc.ErrorIsNil)
	server := newServer(c)
	var opts ssh.Options
	opts.SetPort(server.listener.Addr().(*net.TCPAddr).Port)
	cmd := client.Command("127.0.0.1", testCommand, &opts)
	checkedKey := false
	server.cfg.PublicKeyCallback = func(conn cryptossh.ConnMetadata, pubkey cryptossh.PublicKey) (*cryptossh.Permissions, error) {
		c.Check(pubkey, gc.DeepEquals, key.PublicKey())
		checkedKey = true
		return nil, nil
	}
	go server.run(c)
	out, err := cmd.Output()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(string(out), gc.Equals, "abc value\n")
	c.Assert(checkedKey, jc.IsTrue)
}

// InitUbuntuUser adds the ubuntu user if it doesn't
// already exist, updates its ~/.ssh/authorized_keys,
// and enables passwordless sudo for it.
//
// InitUbuntuUser will initially attempt to login as
// the ubuntu user, and verify that passwordless sudo
// is enabled; only if this is false will there be an
// attempt with the specified login.
//
// authorizedKeys may be empty, in which case the file
// will be created and left empty.
//
// stdin and stdout will be used for remote sudo prompts,
// if the ubuntu user must be created/updated.
func InitUbuntuUser(host, login, authorizedKeys string, stdin io.Reader, stdout io.Writer) error {
	logger.Infof("initialising %q, user %q", host, login)

	// To avoid unnecessary prompting for the specified login,
	// initUbuntuUser will first attempt to ssh to the machine
	// as "ubuntu" with password authentication disabled, and
	// ensure that it can use sudo without a password.
	//
	// Note that we explicitly do not allocate a PTY, so we
	// get a failure if sudo prompts.
	cmd := ssh.Command("[email protected]"+host, []string{"sudo", "-n", "true"}, nil)
	if cmd.Run() == nil {
		logger.Infof("ubuntu user is already initialised")
		return nil
	}

	// Failed to login as ubuntu (or passwordless sudo is not enabled).
	// Use specified login, and execute the initUbuntuScript below.
	if login != "" {
		host = login + "@" + host
	}
	script := fmt.Sprintf(initUbuntuScript, utils.ShQuote(authorizedKeys))
	var options ssh.Options
	options.AllowPasswordAuthentication()
	options.EnablePTY()
	cmd = ssh.Command(host, []string{"sudo", "/bin/bash -c " + utils.ShQuote(script)}, &options)
	var stderr bytes.Buffer
	cmd.Stdin = stdin
	cmd.Stdout = stdout // for sudo prompt
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		if stderr.Len() != 0 {
			err = fmt.Errorf("%v (%v)", err, strings.TrimSpace(stderr.String()))
		}
		return err
	}
	return nil
}

// getSSHOptions configures SSH options based on command line
// arguments and the SSH targets specified.
func (c *SSHCommon) getSSHOptions(enablePty bool, targets ...*resolvedTarget) (*ssh.Options, error) {
	var options ssh.Options

	if c.noHostKeyChecks {
		options.SetStrictHostKeyChecking(ssh.StrictHostChecksNo)
		options.SetKnownHostsFile("/dev/null")
	} else {
		knownHostsPath, err := c.generateKnownHosts(targets)
		if err != nil {
			return nil, errors.Trace(err)
		}

		// There might not be a custom known_hosts file if the SSH
		// targets are specified using arbitrary hostnames or
		// addresses. In this case, the user's personal known_hosts
		// file is used.

		if knownHostsPath != "" {
			// When a known_hosts file has been generated, enforce
			// strict host key checking.
			options.SetStrictHostKeyChecking(ssh.StrictHostChecksYes)
			options.SetKnownHostsFile(knownHostsPath)
		} else {
			// If the user's personal known_hosts is used, also use
			// the user's personal StrictHostKeyChecking preferences.
			options.SetStrictHostKeyChecking(ssh.StrictHostChecksUnset)
		}
	}

	if enablePty {
		options.EnablePTY()
	}

	if c.proxy {
		if err := c.setProxyCommand(&options); err != nil {
			return nil, err
		}
	}

	return &options, nil
}