1/*
2Copyright (c) 2014-2015 VMware, Inc. All Rights Reserved.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8    http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17package vm
18
19import (
20	"context"
21	"encoding/json"
22	"flag"
23	"fmt"
24	"io"
25	"reflect"
26	"regexp"
27	"strconv"
28	"strings"
29
30	"github.com/vmware/govmomi/govc/cli"
31	"github.com/vmware/govmomi/govc/flags"
32	"github.com/vmware/govmomi/object"
33	"github.com/vmware/govmomi/property"
34	"github.com/vmware/govmomi/vim25"
35	"github.com/vmware/govmomi/vim25/mo"
36	"github.com/vmware/govmomi/vim25/types"
37)
38
39type intRange struct {
40	low, high int
41}
42
43var intRangeRegexp = regexp.MustCompile("^([0-9]+)-([0-9]+)$")
44
45func (i *intRange) Set(s string) error {
46	m := intRangeRegexp.FindStringSubmatch(s)
47	if m == nil {
48		return fmt.Errorf("invalid range: %s", s)
49	}
50
51	low, _ := strconv.Atoi(m[1])
52	high, _ := strconv.Atoi(m[2])
53	if low > high {
54		return fmt.Errorf("invalid range: low > high")
55	}
56
57	i.low = low
58	i.high = high
59	return nil
60}
61
62func (i *intRange) String() string {
63	return fmt.Sprintf("%d-%d", i.low, i.high)
64}
65
66type vnc struct {
67	*flags.SearchFlag
68
69	Enable    bool
70	Disable   bool
71	Port      int
72	PortRange intRange
73	Password  string
74}
75
76func init() {
77	cmd := &vnc{}
78	cmd.PortRange.Set("5900-5999")
79	cli.Register("vm.vnc", cmd)
80}
81
82func (cmd *vnc) Register(ctx context.Context, f *flag.FlagSet) {
83	cmd.SearchFlag, ctx = flags.NewSearchFlag(ctx, flags.SearchVirtualMachines)
84	cmd.SearchFlag.Register(ctx, f)
85
86	f.BoolVar(&cmd.Enable, "enable", false, "Enable VNC")
87	f.BoolVar(&cmd.Disable, "disable", false, "Disable VNC")
88	f.IntVar(&cmd.Port, "port", -1, "VNC port (-1 for auto-select)")
89	f.Var(&cmd.PortRange, "port-range", "VNC port auto-select range")
90	f.StringVar(&cmd.Password, "password", "", "VNC password")
91}
92
93func (cmd *vnc) Process(ctx context.Context) error {
94	if err := cmd.SearchFlag.Process(ctx); err != nil {
95		return err
96	}
97	// Either may be true or none may be true.
98	if cmd.Enable && cmd.Disable {
99		return flag.ErrHelp
100	}
101
102	return nil
103}
104
105func (cmd *vnc) Usage() string {
106	return "VM..."
107}
108
109func (cmd *vnc) Description() string {
110	return `Enable or disable VNC for VM.
111
112Port numbers are automatically chosen if not specified.
113
114If neither -enable or -disable is specified, the current state is returned.
115
116Examples:
117  govc vm.vnc -enable -password 1234 $vm | awk '{print $2}' | xargs open`
118}
119
120func (cmd *vnc) Run(ctx context.Context, f *flag.FlagSet) error {
121	vms, err := cmd.loadVMs(f.Args())
122	if err != nil {
123		return err
124	}
125
126	// Actuate settings in VMs
127	for _, vm := range vms {
128		switch {
129		case cmd.Enable:
130			vm.enable(cmd.Port, cmd.Password)
131		case cmd.Disable:
132			vm.disable()
133		}
134	}
135
136	// Reconfigure VMs to reflect updates
137	for _, vm := range vms {
138		err = vm.reconfigure()
139		if err != nil {
140			return err
141		}
142	}
143
144	return cmd.WriteResult(vncResult(vms))
145}
146
147func (cmd *vnc) loadVMs(args []string) ([]*vncVM, error) {
148	c, err := cmd.Client()
149	if err != nil {
150		return nil, err
151	}
152
153	vms, err := cmd.VirtualMachines(args)
154	if err != nil {
155		return nil, err
156	}
157
158	var vncVMs []*vncVM
159	for _, vm := range vms {
160		v, err := newVNCVM(c, vm)
161		if err != nil {
162			return nil, err
163		}
164		vncVMs = append(vncVMs, v)
165	}
166
167	// Assign vncHosts to vncVMs
168	hosts := make(map[string]*vncHost)
169	for _, vm := range vncVMs {
170		if h, ok := hosts[vm.hostReference().Value]; ok {
171			vm.host = h
172			continue
173		}
174
175		hs := object.NewHostSystem(c, vm.hostReference())
176		h, err := newVNCHost(c, hs, cmd.PortRange.low, cmd.PortRange.high)
177		if err != nil {
178			return nil, err
179		}
180
181		hosts[vm.hostReference().Value] = h
182		vm.host = h
183	}
184
185	return vncVMs, nil
186}
187
188type vncVM struct {
189	c    *vim25.Client
190	vm   *object.VirtualMachine
191	mvm  mo.VirtualMachine
192	host *vncHost
193
194	curOptions vncOptions
195	newOptions vncOptions
196}
197
198func newVNCVM(c *vim25.Client, vm *object.VirtualMachine) (*vncVM, error) {
199	v := &vncVM{
200		c:  c,
201		vm: vm,
202	}
203
204	virtualMachineProperties := []string{
205		"name",
206		"config.extraConfig",
207		"runtime.host",
208	}
209
210	pc := property.DefaultCollector(c)
211	ctx := context.TODO()
212	err := pc.RetrieveOne(ctx, vm.Reference(), virtualMachineProperties, &v.mvm)
213	if err != nil {
214		return nil, err
215	}
216
217	v.curOptions = vncOptionsFromExtraConfig(v.mvm.Config.ExtraConfig)
218	v.newOptions = vncOptionsFromExtraConfig(v.mvm.Config.ExtraConfig)
219
220	return v, nil
221}
222
223func (v *vncVM) hostReference() types.ManagedObjectReference {
224	return *v.mvm.Runtime.Host
225}
226
227func (v *vncVM) enable(port int, password string) error {
228	v.newOptions["enabled"] = "true"
229	v.newOptions["port"] = fmt.Sprintf("%d", port)
230	v.newOptions["password"] = password
231
232	// Find port if auto-select
233	if port == -1 {
234		// Reuse port if If VM already has a port, reuse it.
235		// Otherwise, find unused VNC port on host.
236		if p, ok := v.curOptions["port"]; ok && p != "" {
237			v.newOptions["port"] = p
238		} else {
239			port, err := v.host.popUnusedPort()
240			if err != nil {
241				return err
242			}
243			v.newOptions["port"] = fmt.Sprintf("%d", port)
244		}
245	}
246	return nil
247}
248
249func (v *vncVM) disable() error {
250	v.newOptions["enabled"] = "false"
251	v.newOptions["port"] = ""
252	v.newOptions["password"] = ""
253	return nil
254}
255
256func (v *vncVM) reconfigure() error {
257	if reflect.DeepEqual(v.curOptions, v.newOptions) {
258		// No changes to settings
259		return nil
260	}
261
262	spec := types.VirtualMachineConfigSpec{
263		ExtraConfig: v.newOptions.ToExtraConfig(),
264	}
265
266	ctx := context.TODO()
267	task, err := v.vm.Reconfigure(ctx, spec)
268	if err != nil {
269		return err
270	}
271
272	return task.Wait(ctx)
273}
274
275func (v *vncVM) uri() (string, error) {
276	ip, err := v.host.managementIP()
277	if err != nil {
278		return "", err
279	}
280
281	uri := fmt.Sprintf("vnc://:%s@%s:%s",
282		v.newOptions["password"],
283		ip,
284		v.newOptions["port"])
285
286	return uri, nil
287}
288
289func (v *vncVM) write(w io.Writer) error {
290	if strings.EqualFold(v.newOptions["enabled"], "true") {
291		uri, err := v.uri()
292		if err != nil {
293			return err
294		}
295		fmt.Printf("%s: %s\n", v.mvm.Name, uri)
296	} else {
297		fmt.Printf("%s: disabled\n", v.mvm.Name)
298	}
299	return nil
300}
301
302type vncHost struct {
303	c     *vim25.Client
304	host  *object.HostSystem
305	ports map[int]struct{}
306	ip    string // This field is populated by `managementIP`
307}
308
309func newVNCHost(c *vim25.Client, host *object.HostSystem, low, high int) (*vncHost, error) {
310	ports := make(map[int]struct{})
311	for i := low; i <= high; i++ {
312		ports[i] = struct{}{}
313	}
314
315	used, err := loadUsedPorts(c, host.Reference())
316	if err != nil {
317		return nil, err
318	}
319
320	// Remove used ports from range
321	for _, u := range used {
322		delete(ports, u)
323	}
324
325	h := &vncHost{
326		c:     c,
327		host:  host,
328		ports: ports,
329	}
330
331	return h, nil
332}
333
334func loadUsedPorts(c *vim25.Client, host types.ManagedObjectReference) ([]int, error) {
335	ctx := context.TODO()
336	ospec := types.ObjectSpec{
337		Obj: host,
338		SelectSet: []types.BaseSelectionSpec{
339			&types.TraversalSpec{
340				Type: "HostSystem",
341				Path: "vm",
342				Skip: types.NewBool(false),
343			},
344		},
345		Skip: types.NewBool(false),
346	}
347
348	pspec := types.PropertySpec{
349		Type:    "VirtualMachine",
350		PathSet: []string{"config.extraConfig"},
351	}
352
353	req := types.RetrieveProperties{
354		This: c.ServiceContent.PropertyCollector,
355		SpecSet: []types.PropertyFilterSpec{
356			{
357				ObjectSet: []types.ObjectSpec{ospec},
358				PropSet:   []types.PropertySpec{pspec},
359			},
360		},
361	}
362
363	var vms []mo.VirtualMachine
364	err := mo.RetrievePropertiesForRequest(ctx, c, req, &vms)
365	if err != nil {
366		return nil, err
367	}
368
369	var ports []int
370	for _, vm := range vms {
371		if vm.Config == nil || vm.Config.ExtraConfig == nil {
372			continue
373		}
374
375		options := vncOptionsFromExtraConfig(vm.Config.ExtraConfig)
376		if ps, ok := options["port"]; ok && ps != "" {
377			pi, err := strconv.Atoi(ps)
378			if err == nil {
379				ports = append(ports, pi)
380			}
381		}
382	}
383
384	return ports, nil
385}
386
387func (h *vncHost) popUnusedPort() (int, error) {
388	if len(h.ports) == 0 {
389		return 0, fmt.Errorf("no unused ports in range")
390	}
391
392	// Return first port we get when iterating
393	var port int
394	for port = range h.ports {
395		break
396	}
397	delete(h.ports, port)
398	return port, nil
399}
400
401func (h *vncHost) managementIP() (string, error) {
402	ctx := context.TODO()
403	if h.ip != "" {
404		return h.ip, nil
405	}
406
407	ips, err := h.host.ManagementIPs(ctx)
408	if err != nil {
409		return "", err
410	}
411
412	if len(ips) > 0 {
413		h.ip = ips[0].String()
414	} else {
415		h.ip = "<unknown>"
416	}
417
418	return h.ip, nil
419}
420
421type vncResult []*vncVM
422
423func (vms vncResult) MarshalJSON() ([]byte, error) {
424	out := make(map[string]string)
425	for _, vm := range vms {
426		uri, err := vm.uri()
427		if err != nil {
428			return nil, err
429		}
430		out[vm.mvm.Name] = uri
431	}
432	return json.Marshal(out)
433}
434
435func (vms vncResult) Write(w io.Writer) error {
436	for _, vm := range vms {
437		err := vm.write(w)
438		if err != nil {
439			return err
440		}
441	}
442
443	return nil
444}
445
446type vncOptions map[string]string
447
448var vncPrefix = "RemoteDisplay.vnc."
449
450func vncOptionsFromExtraConfig(ov []types.BaseOptionValue) vncOptions {
451	vo := make(vncOptions)
452	for _, b := range ov {
453		o := b.GetOptionValue()
454		if strings.HasPrefix(o.Key, vncPrefix) {
455			key := o.Key[len(vncPrefix):]
456			if key != "key" {
457				vo[key] = o.Value.(string)
458			}
459		}
460	}
461	return vo
462}
463
464func (vo vncOptions) ToExtraConfig() []types.BaseOptionValue {
465	ov := make([]types.BaseOptionValue, 0, 0)
466	for k, v := range vo {
467		key := vncPrefix + k
468		value := v
469
470		o := types.OptionValue{
471			Key:   key,
472			Value: &value, // Pass pointer to avoid omitempty
473		}
474
475		ov = append(ov, &o)
476	}
477
478	// Don't know how to deal with the key option, set it to be empty...
479	o := types.OptionValue{
480		Key:   vncPrefix + "key",
481		Value: new(string), // Pass pointer to avoid omitempty
482	}
483
484	ov = append(ov, &o)
485
486	return ov
487}
488