// Copyright 2015 The Vanadium 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 (
	"bytes"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"v.io/x/lib/cmdline"
)

var cmdNode = &cmdline.Command{
	Name:     "node",
	Short:    "Manage GCE nodes",
	Long:     "Manage GCE nodes.",
	Children: []*cmdline.Command{cmdNodeAuthorize, cmdNodeDeauthorize, cmdNodeCreate, cmdNodeDelete},
}

var cmdNodeAuthorize = &cmdline.Command{
	Runner: cmdline.RunnerFunc(runNodeAuthorize),
	Name:   "authorize",
	Short:  "Authorize a user to login to a GCE node",
	Long: `
Authorizes a user to login to a GCE node (possibly as other user). For
instance, this mechanism is used to give Jenkins slave nodes access to
the GCE mirror of Vanadium repositories.
`,
	ArgsName: "<userA>@<hostA> [<userB>@]<hostB>",
	ArgsLong: `
<userA>@<hostA> [<userB>@]<hostB> authorizes userA to log into GCE
node hostB from GCE node hostA as user userB. The default value for
userB is userA.
`,
}

var cmdNodeDeauthorize = &cmdline.Command{
	Runner: cmdline.RunnerFunc(runNodeDeauthorize),
	Name:   "deauthorize",
	Short:  "Deauthorize a user to login to a GCE node",
	Long: `
Deuthorizes a user to login to a GCE node (possibly as other
user). For instance, this mechanism is used to revoke access of give
Jenkins slave nodes to the GCE mirror of Vanadium repositories.
`,
	ArgsName: "<userA>@<hostA> [<userB>@]<hostB>",
	ArgsLong: `
<userA>@<hostA> [<userB>@]<hostB> deauthorizes userA to log into GCE
node hostB from GCE node hostA as user userB. The default value for
userB is userA.
`,
}

func parseUserAndHost(args []string) (string, string, string, string, error) {
	if got, want := len(args), 2; got != want {
		return "", "", "", "", fmt.Errorf("unexpected number of arguments: got %v, want %d", got, want)
	}

	parseFn := func(s string) (string, string, error) {
		tokens := strings.Split(s, "@")
		switch len(tokens) {
		case 1:
			return "", tokens[0], nil
		case 2:
			return tokens[0], tokens[1], nil
		default:
			return "", "", fmt.Errorf("unexpected length of %v: expected at most %d", tokens, 2)
		}
	}

	userA, hostA, err := parseFn(args[0])
	if err != nil {
		return "", "", "", "", err
	}
	if userA == "" {
		return "", "", "", "", fmt.Errorf("failed to parse user: %v", args[0])
	}
	userB, hostB, err := parseFn(args[1])
	if err != nil {
		return "", "", "", "", err
	}
	if userB == "" {
		userB = userA
	}

	return userA, hostA, userB, hostB, nil
}

// TODO(jsimsa): Add command-line flags for specifying the name of the
// SSH key file to use and whether to create one if it does not
// exist.
func runNodeAuthorize(env *cmdline.Env, args []string) error {
	userA, hostA, userB, hostB, err := parseUserAndHost(args)
	if err != nil {
		return env.UsageErrorf("%v", err)
	}

	// Copy the public SSH key for <userA> from <hostA> to the local
	// machine.
	ctx := newContext(env)
	s := ctx.NewSeq()
	tmpDir, err := s.TempDir("", "")
	if err != nil {
		return fmt.Errorf("TempDir() failed: %v", err)
	}
	defer ctx.NewSeq().RemoveAll(tmpDir)
	allNodes, err := listAll(ctx)
	if err != nil {
		return err
	}
	nodeA, err := allNodes.MatchNames(hostA)
	if err != nil {
		return err
	}
	if err := nodeA.RunCopy(ctx, []string{fmt.Sprintf(":/home/%v/.ssh/id_rsa.pub", userA)}, tmpDir); err != nil {
		return err
	}

	// Append the key to the set of authorized keys of <userB> on
	// <hostB>.
	sshKeyFile := filepath.Join(tmpDir, "id_rsa.pub")
	bytes, err := s.ReadFile(sshKeyFile)
	if err != nil {
		return fmt.Errorf("ReadFile(%v) failed: %v", sshKeyFile, err)
	}
	nodeB, err := allNodes.MatchNames(hostB)
	if err != nil {
		return err
	}
	echoCmd := []string{"echo", strings.TrimSpace(string(bytes)), ">>", fmt.Sprintf("/home/%v/.ssh/authorized_keys", userB)}
	if err := nodeB.RunCommand(ctx, userB, echoCmd); err != nil {
		return err
	}
	return nil
}

func runNodeDeauthorize(env *cmdline.Env, args []string) error {
	userA, hostA, userB, hostB, err := parseUserAndHost(args)
	if err != nil {
		return env.UsageErrorf("%v", err)
	}

	// Remove all keys for <userA>@<hostA> from the set of authorized
	// keys of <userB> on <hostB>.
	ctx := newContext(env)
	allNodes, err := listAll(ctx)
	if err != nil {
		return err
	}
	nodeB, err := allNodes.MatchNames(hostB)
	if err != nil {
		return err
	}
	authorizedKeysFile := fmt.Sprintf("/home/%v/.ssh/authorized_keys", userB)
	tmpKeysFile := authorizedKeysFile + ".tmp"
	grepCmd := []string{"grep", "-v", fmt.Sprintf("%v@%v", userA, hostA), authorizedKeysFile, ">", tmpKeysFile}
	if err := nodeB.RunCommand(ctx, userB, grepCmd); err != nil {
		return err
	}
	moveCmd := []string{"mv", tmpKeysFile, authorizedKeysFile}
	if err := nodeB.RunCommand(ctx, userB, moveCmd); err != nil {
		return err
	}

	return nil
}

var cmdNodeCreate = &cmdline.Command{
	Runner: cmdline.RunnerFunc(runNodeCreate),
	Name:   "create",
	Short:  "Create GCE nodes",
	Long: `
Create GCE nodes. Runs 'gcloud compute instances create'.
`,
	ArgsName: "<names>",
	ArgsLong: "<names> is a list of names identifying nodes to be created.",
}

var cmdNodeDelete = &cmdline.Command{
	Runner: cmdline.RunnerFunc(runNodeDelete),
	Name:   "delete",
	Short:  "Delete GCE nodes",
	Long: `
Delete GCE nodes. Runs 'gcloud compute instances delete'.
`,
	ArgsName: "<names>",
	ArgsLong: "<names> is a list of names identifying nodes to be deleted.",
}

func runNodeCreate(env *cmdline.Env, args []string) error {
	ctx := newContext(env)

	// Create the GCE node(s).
	createArgs := []string{
		"compute",
		"--project", *flagProject,
		"instances",
		"create",
	}
	createArgs = append(createArgs, args...)
	createArgs = append(createArgs,
		"--boot-disk-size", flagBootDiskSize,
		"--image", flagImage,
		"--machine-type", flagMachineType,
		"--zone", flagZone,
		"--scopes", flagScopes,
	)
	if err := ctx.NewSeq().Last("gcloud", createArgs...); err != nil {
		return err
	}

	// Create in-memory representation of node information.
	allNodes, err := listAll(ctx)
	if err != nil {
		return err
	}
	nodes, err := allNodes.MatchNames(strings.Join(args, ","))
	if err != nil {
		return err
	}

	// Wait for the SSH server on all nodes to start up.
	const numRetries = 10
	const retryPeriod = 5 * time.Second
	ready := false
	for i := 0; i < numRetries; i++ {
		if err := nodes.RunCommand(ctx, *flagUser, []string{"echo"}); err != nil {
			fmt.Fprintf(ctx.Stdout(), "attempt #%d to connect failed, will try again later\n", i+1)
			time.Sleep(retryPeriod)
			continue
		}
		ready = true
		break
	}
	if !ready {
		return fmt.Errorf("timed out waiting for nodes to start")
	}

	// Execute the setup script.
	if flagSetupScript != "" {
		if err := nodes.RunCopyAndRun(ctx, *flagUser, []string{flagSetupScript}, nil, ""); err != nil {
			return err
		}
	}

	return nil
}

func runNodeDelete(env *cmdline.Env, args []string) error {
	ctx := newContext(env)
	// Delete the GCE node(s).
	var in bytes.Buffer
	in.WriteString("Y\n") // answers the [Y/n] prompt
	deleteArgs := []string{
		"compute",
		"--project", *flagProject,
		"instances",
		"delete",
	}
	deleteArgs = append(deleteArgs, args...)
	deleteArgs = append(deleteArgs, "--zone", flagZone)
	return ctx.NewSeq().Read(&in).Last("gcloud", deleteArgs...)
}