vnc.go

/*
Copyright (c) 2014-2015 VMware, Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vm

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/vmware/govmomi/govc/cli"
	"github.com/vmware/govmomi/govc/flags"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/property"
	"github.com/vmware/govmomi/vim25"
	"github.com/vmware/govmomi/vim25/mo"
	"github.com/vmware/govmomi/vim25/types"
	"golang.org/x/net/context"
)

type intRange struct {
	low, high int
}

var intRangeRegexp = regexp.MustCompile("^([0-9]+)-([0-9]+)$")

func (i *intRange) Set(s string) error {
	m := intRangeRegexp.FindStringSubmatch(s)
	if m == nil {
		return fmt.Errorf("invalid range: %s", s)
	}

	low, _ := strconv.Atoi(m[1])
	high, _ := strconv.Atoi(m[2])
	if low > high {
		return fmt.Errorf("invalid range: low > high")
	}

	i.low = low
	i.high = high
	return nil
}

func (i *intRange) String() string {
	return fmt.Sprintf("%d-%d", i.low, i.high)
}

type vnc struct {
	*flags.SearchFlag

	Enable    bool
	Disable   bool
	Port      int
	PortRange intRange
	Password  string
}

func init() {
	cmd := &vnc{}
	cmd.PortRange.Set("5900-5999")
	cli.Register("vm.vnc", cmd)
}

func (cmd *vnc) Register(ctx context.Context, f *flag.FlagSet) {
	cmd.SearchFlag, ctx = flags.NewSearchFlag(ctx, flags.SearchVirtualMachines)
	cmd.SearchFlag.Register(ctx, f)

	f.BoolVar(&cmd.Enable, "enable", false, "Enable VNC")
	f.BoolVar(&cmd.Disable, "disable", false, "Disable VNC")
	f.IntVar(&cmd.Port, "port", -1, "VNC port (-1 for auto-select)")
	f.Var(&cmd.PortRange, "port-range", "VNC port auto-select range")
	f.StringVar(&cmd.Password, "password", "", "VNC password")
}

func (cmd *vnc) Process(ctx context.Context) error {
	if err := cmd.SearchFlag.Process(ctx); err != nil {
		return err
	}
	// Either may be true or none may be true.
	if cmd.Enable && cmd.Disable {
		return flag.ErrHelp
	}

	return nil
}

func (cmd *vnc) Usage() string {
	return "VM..."
}

func (cmd *vnc) Description() string {
	return `VNC controls for VM(s).

Port numbers are automatically chosen from a range if not specified.

If neither -enable or -disable is specified, the current state is returned.`
}

func (cmd *vnc) Run(ctx context.Context, f *flag.FlagSet) error {
	vms, err := cmd.loadVMs(f.Args())
	if err != nil {
		return err
	}

	// Actuate settings in VMs
	for _, vm := range vms {
		switch {
		case cmd.Enable:
			vm.enable(cmd.Port, cmd.Password)
		case cmd.Disable:
			vm.disable()
		}
	}

	// Reconfigure VMs to reflect updates
	for _, vm := range vms {
		err = vm.reconfigure()
		if err != nil {
			return err
		}
	}

	return cmd.WriteResult(vncResult(vms))
}

func (cmd *vnc) loadVMs(args []string) ([]*vncVM, error) {
	c, err := cmd.Client()
	if err != nil {
		return nil, err
	}

	vms, err := cmd.VirtualMachines(args)
	if err != nil {
		return nil, err
	}

	var vncVMs []*vncVM
	for _, vm := range vms {
		v, err := newVNCVM(c, vm)
		if err != nil {
			return nil, err
		}
		vncVMs = append(vncVMs, v)
	}

	// Assign vncHosts to vncVMs
	hosts := make(map[string]*vncHost)
	for _, vm := range vncVMs {
		if h, ok := hosts[vm.hostReference().Value]; ok {
			vm.host = h
			continue
		}

		hs := object.NewHostSystem(c, vm.hostReference())
		h, err := newVNCHost(c, hs, cmd.PortRange.low, cmd.PortRange.high)
		if err != nil {
			return nil, err
		}

		hosts[vm.hostReference().Value] = h
		vm.host = h
	}

	return vncVMs, nil
}

type vncVM struct {
	c    *vim25.Client
	vm   *object.VirtualMachine
	mvm  mo.VirtualMachine
	host *vncHost

	curOptions vncOptions
	newOptions vncOptions
}

func newVNCVM(c *vim25.Client, vm *object.VirtualMachine) (*vncVM, error) {
	v := &vncVM{
		c:  c,
		vm: vm,
	}

	virtualMachineProperties := []string{
		"name",
		"config.extraConfig",
		"runtime.host",
	}

	pc := property.DefaultCollector(c)
	err := pc.RetrieveOne(context.TODO(), vm.Reference(), virtualMachineProperties, &v.mvm)
	if err != nil {
		return nil, err
	}

	v.curOptions = vncOptionsFromExtraConfig(v.mvm.Config.ExtraConfig)
	v.newOptions = vncOptionsFromExtraConfig(v.mvm.Config.ExtraConfig)

	return v, nil
}

func (v *vncVM) hostReference() types.ManagedObjectReference {
	return *v.mvm.Runtime.Host
}

func (v *vncVM) enable(port int, password string) error {
	v.newOptions["enabled"] = "true"
	v.newOptions["port"] = fmt.Sprintf("%d", port)
	v.newOptions["password"] = password

	// Find port if auto-select
	if port == -1 {
		// Reuse port if If VM already has a port, reuse it.
		// Otherwise, find unused VNC port on host.
		if p, ok := v.curOptions["port"]; ok && p != "" {
			v.newOptions["port"] = p
		} else {
			port, err := v.host.popUnusedPort()
			if err != nil {
				return err
			}
			v.newOptions["port"] = fmt.Sprintf("%d", port)
		}
	}
	return nil
}

func (v *vncVM) disable() error {
	v.newOptions["enabled"] = "false"
	v.newOptions["port"] = ""
	v.newOptions["password"] = ""
	return nil
}

func (v *vncVM) reconfigure() error {
	if reflect.DeepEqual(v.curOptions, v.newOptions) {
		// No changes to settings
		return nil
	}

	spec := types.VirtualMachineConfigSpec{
		ExtraConfig: v.newOptions.ToExtraConfig(),
	}

	task, err := v.vm.Reconfigure(context.TODO(), spec)
	if err != nil {
		return err
	}

	return task.Wait(context.TODO())
}

func (v *vncVM) uri() (string, error) {
	ip, err := v.host.managementIP()
	if err != nil {
		return "", err
	}

	uri := fmt.Sprintf("vnc://:%s@%s:%s",
		v.newOptions["password"],
		ip,
		v.newOptions["port"])

	return uri, nil
}

func (v *vncVM) write(w io.Writer) error {
	if v.newOptions["enabled"] == "true" {
		uri, err := v.uri()
		if err != nil {
			return err
		}
		fmt.Printf("%s: %s\n", v.mvm.Name, uri)
	} else {
		fmt.Printf("%s: disabled\n", v.mvm.Name)
	}
	return nil
}

type vncHost struct {
	c     *vim25.Client
	host  *object.HostSystem
	ports map[int]struct{}
	ip    string // This field is populated by `managementIP`
}

func newVNCHost(c *vim25.Client, host *object.HostSystem, low, high int) (*vncHost, error) {
	ports := make(map[int]struct{})
	for i := low; i <= high; i++ {
		ports[i] = struct{}{}
	}

	used, err := loadUsedPorts(c, host.Reference())
	if err != nil {
		return nil, err
	}

	// Remove used ports from range
	for _, u := range used {
		delete(ports, u)
	}

	h := &vncHost{
		c:     c,
		host:  host,
		ports: ports,
	}

	return h, nil
}

func loadUsedPorts(c *vim25.Client, host types.ManagedObjectReference) ([]int, error) {
	ospec := types.ObjectSpec{
		Obj: host,
		SelectSet: []types.BaseSelectionSpec{
			&types.TraversalSpec{
				Type: "HostSystem",
				Path: "vm",
				Skip: types.NewBool(false),
			},
		},
		Skip: types.NewBool(false),
	}

	pspec := types.PropertySpec{
		Type:    "VirtualMachine",
		PathSet: []string{"config.extraConfig"},
	}

	req := types.RetrieveProperties{
		This: c.ServiceContent.PropertyCollector,
		SpecSet: []types.PropertyFilterSpec{
			{
				ObjectSet: []types.ObjectSpec{ospec},
				PropSet:   []types.PropertySpec{pspec},
			},
		},
	}

	var vms []mo.VirtualMachine
	err := mo.RetrievePropertiesForRequest(context.TODO(), c, req, &vms)
	if err != nil {
		return nil, err
	}

	var ports []int
	for _, vm := range vms {
		if vm.Config == nil || vm.Config.ExtraConfig == nil {
			continue
		}

		options := vncOptionsFromExtraConfig(vm.Config.ExtraConfig)
		if ps, ok := options["port"]; ok && ps != "" {
			pi, err := strconv.Atoi(ps)
			if err == nil {
				ports = append(ports, pi)
			}
		}
	}

	return ports, nil
}

func (h *vncHost) popUnusedPort() (int, error) {
	if len(h.ports) == 0 {
		return 0, fmt.Errorf("no unused ports in range")
	}

	// Return first port we get when iterating
	var port int
	for port, _ = range h.ports {
		break
	}
	delete(h.ports, port)
	return port, nil
}

func (h *vncHost) managementIP() (string, error) {
	if h.ip != "" {
		return h.ip, nil
	}

	ips, err := h.host.ManagementIPs(context.TODO())
	if err != nil {
		return "", err
	}

	if len(ips) > 0 {
		h.ip = ips[0].String()
	} else {
		h.ip = "<unknown>"
	}

	return h.ip, nil
}

type vncResult []*vncVM

func (vms vncResult) MarshalJSON() ([]byte, error) {
	out := make(map[string]string)
	for _, vm := range vms {
		uri, err := vm.uri()
		if err != nil {
			return nil, err
		}
		out[vm.mvm.Name] = uri
	}
	return json.Marshal(out)
}

func (vms vncResult) Write(w io.Writer) error {
	for _, vm := range vms {
		err := vm.write(w)
		if err != nil {
			return err
		}
	}

	return nil
}

type vncOptions map[string]string

var vncPrefix = "RemoteDisplay.vnc."

func vncOptionsFromExtraConfig(ov []types.BaseOptionValue) vncOptions {
	vo := make(vncOptions)
	for _, b := range ov {
		o := b.GetOptionValue()
		if strings.HasPrefix(o.Key, vncPrefix) {
			key := o.Key[len(vncPrefix):]
			if key != "key" {
				vo[key] = o.Value.(string)
			}
		}
	}
	return vo
}

func (vo vncOptions) ToExtraConfig() []types.BaseOptionValue {
	ov := make([]types.BaseOptionValue, 0, 0)
	for k, v := range vo {
		key := vncPrefix + k
		value := v

		o := types.OptionValue{
			Key:   key,
			Value: &value, // Pass pointer to avoid omitempty
		}

		ov = append(ov, &o)
	}

	// Don't know how to deal with the key option, set it to be empty...
	o := types.OptionValue{
		Key:   vncPrefix + "key",
		Value: new(string), // Pass pointer to avoid omitempty
	}

	ov = append(ov, &o)

	return ov
}