1/*
2Copyright (c) 2014-2016 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 flags
18
19import (
20	"context"
21	"errors"
22	"flag"
23	"fmt"
24	"strings"
25
26	"github.com/vmware/govmomi/find"
27	"github.com/vmware/govmomi/object"
28	"github.com/vmware/govmomi/vim25"
29	"github.com/vmware/govmomi/vim25/soap"
30	"github.com/vmware/govmomi/vim25/types"
31)
32
33const (
34	SearchVirtualMachines = iota + 1
35	SearchHosts
36	SearchVirtualApps
37)
38
39type SearchFlag struct {
40	common
41
42	*ClientFlag
43	*DatacenterFlag
44
45	t      int
46	entity string
47
48	byDatastorePath string
49	byDNSName       string
50	byInventoryPath string
51	byIP            string
52	byUUID          string
53
54	isset bool
55}
56
57var searchFlagKey = flagKey("search")
58
59func NewSearchFlag(ctx context.Context, t int) (*SearchFlag, context.Context) {
60	if v := ctx.Value(searchFlagKey); v != nil {
61		return v.(*SearchFlag), ctx
62	}
63
64	v := &SearchFlag{
65		t: t,
66	}
67
68	v.ClientFlag, ctx = NewClientFlag(ctx)
69	v.DatacenterFlag, ctx = NewDatacenterFlag(ctx)
70
71	switch t {
72	case SearchVirtualMachines:
73		v.entity = "VM"
74	case SearchHosts:
75		v.entity = "host"
76	case SearchVirtualApps:
77		v.entity = "vapp"
78	default:
79		panic("invalid search type")
80	}
81
82	ctx = context.WithValue(ctx, searchFlagKey, v)
83	return v, ctx
84}
85
86func (flag *SearchFlag) Register(ctx context.Context, fs *flag.FlagSet) {
87	flag.RegisterOnce(func() {
88		flag.ClientFlag.Register(ctx, fs)
89		flag.DatacenterFlag.Register(ctx, fs)
90
91		register := func(v *string, f string, d string) {
92			f = fmt.Sprintf("%s.%s", strings.ToLower(flag.entity), f)
93			d = fmt.Sprintf(d, flag.entity)
94			fs.StringVar(v, f, "", d)
95		}
96
97		switch flag.t {
98		case SearchVirtualMachines:
99			register(&flag.byDatastorePath, "path", "Find %s by path to .vmx file")
100		}
101
102		switch flag.t {
103		case SearchVirtualMachines, SearchHosts:
104			register(&flag.byDNSName, "dns", "Find %s by FQDN")
105			register(&flag.byIP, "ip", "Find %s by IP address")
106			register(&flag.byUUID, "uuid", "Find %s by UUID")
107		}
108
109		register(&flag.byInventoryPath, "ipath", "Find %s by inventory path")
110	})
111}
112
113func (flag *SearchFlag) Process(ctx context.Context) error {
114	return flag.ProcessOnce(func() error {
115		if err := flag.ClientFlag.Process(ctx); err != nil {
116			return err
117		}
118		if err := flag.DatacenterFlag.Process(ctx); err != nil {
119			return err
120		}
121
122		flags := []string{
123			flag.byDatastorePath,
124			flag.byDNSName,
125			flag.byInventoryPath,
126			flag.byIP,
127			flag.byUUID,
128		}
129
130		flag.isset = false
131		for _, f := range flags {
132			if f != "" {
133				if flag.isset {
134					return errors.New("cannot use more than one search flag")
135				}
136				flag.isset = true
137			}
138		}
139
140		return nil
141	})
142}
143
144func (flag *SearchFlag) IsSet() bool {
145	return flag.isset
146}
147
148func (flag *SearchFlag) searchIndex(c *vim25.Client) *object.SearchIndex {
149	return object.NewSearchIndex(c)
150}
151
152func (flag *SearchFlag) searchByDatastorePath(c *vim25.Client, dc *object.Datacenter) (object.Reference, error) {
153	ctx := context.TODO()
154	switch flag.t {
155	case SearchVirtualMachines:
156		return flag.searchIndex(c).FindByDatastorePath(ctx, dc, flag.byDatastorePath)
157	default:
158		panic("unsupported type")
159	}
160}
161
162func (flag *SearchFlag) searchByDNSName(c *vim25.Client, dc *object.Datacenter) (object.Reference, error) {
163	ctx := context.TODO()
164	switch flag.t {
165	case SearchVirtualMachines:
166		return flag.searchIndex(c).FindByDnsName(ctx, dc, flag.byDNSName, true)
167	case SearchHosts:
168		return flag.searchIndex(c).FindByDnsName(ctx, dc, flag.byDNSName, false)
169	default:
170		panic("unsupported type")
171	}
172}
173
174func (flag *SearchFlag) searchByInventoryPath(c *vim25.Client, dc *object.Datacenter) (object.Reference, error) {
175	// TODO(PN): The datacenter flag should not be set because it is ignored.
176	ctx := context.TODO()
177	return flag.searchIndex(c).FindByInventoryPath(ctx, flag.byInventoryPath)
178}
179
180func (flag *SearchFlag) searchByIP(c *vim25.Client, dc *object.Datacenter) (object.Reference, error) {
181	ctx := context.TODO()
182	switch flag.t {
183	case SearchVirtualMachines:
184		return flag.searchIndex(c).FindByIp(ctx, dc, flag.byIP, true)
185	case SearchHosts:
186		return flag.searchIndex(c).FindByIp(ctx, dc, flag.byIP, false)
187	default:
188		panic("unsupported type")
189	}
190}
191
192func (flag *SearchFlag) searchByUUID(c *vim25.Client, dc *object.Datacenter) (object.Reference, error) {
193	ctx := context.TODO()
194	isVM := false
195	switch flag.t {
196	case SearchVirtualMachines:
197		isVM = true
198	case SearchHosts:
199	default:
200		panic("unsupported type")
201	}
202
203	var ref object.Reference
204	var err error
205
206	for _, iu := range []*bool{nil, types.NewBool(true)} {
207		ref, err = flag.searchIndex(c).FindByUuid(ctx, dc, flag.byUUID, isVM, iu)
208		if err != nil {
209			if soap.IsSoapFault(err) {
210				fault := soap.ToSoapFault(err).VimFault()
211				if _, ok := fault.(types.InvalidArgument); ok {
212					continue
213				}
214			}
215			return nil, err
216		}
217		if ref != nil {
218			break
219		}
220	}
221
222	return ref, nil
223}
224
225func (flag *SearchFlag) search() (object.Reference, error) {
226	ctx := context.TODO()
227	var ref object.Reference
228	var err error
229
230	c, err := flag.Client()
231	if err != nil {
232		return nil, err
233	}
234
235	dc, err := flag.Datacenter()
236	if err != nil {
237		return nil, err
238	}
239
240	switch {
241	case flag.byDatastorePath != "":
242		ref, err = flag.searchByDatastorePath(c, dc)
243	case flag.byDNSName != "":
244		ref, err = flag.searchByDNSName(c, dc)
245	case flag.byInventoryPath != "":
246		ref, err = flag.searchByInventoryPath(c, dc)
247	case flag.byIP != "":
248		ref, err = flag.searchByIP(c, dc)
249	case flag.byUUID != "":
250		ref, err = flag.searchByUUID(c, dc)
251	default:
252		err = errors.New("no search flag specified")
253	}
254
255	if err != nil {
256		return nil, err
257	}
258
259	if ref == nil {
260		return nil, fmt.Errorf("no such %s", flag.entity)
261	}
262
263	// set the InventoryPath field
264	finder, err := flag.Finder()
265	if err != nil {
266		return nil, err
267	}
268	ref, err = finder.ObjectReference(ctx, ref.Reference())
269	if err != nil {
270		return nil, err
271	}
272
273	return ref, nil
274}
275
276func (flag *SearchFlag) VirtualMachine() (*object.VirtualMachine, error) {
277	ref, err := flag.search()
278	if err != nil {
279		return nil, err
280	}
281
282	vm, ok := ref.(*object.VirtualMachine)
283	if !ok {
284		return nil, fmt.Errorf("expected VirtualMachine entity, got %s", ref.Reference().Type)
285	}
286
287	return vm, nil
288}
289
290func (flag *SearchFlag) VirtualMachines(args []string) ([]*object.VirtualMachine, error) {
291	ctx := context.TODO()
292	var out []*object.VirtualMachine
293
294	if flag.IsSet() {
295		vm, err := flag.VirtualMachine()
296		if err != nil {
297			return nil, err
298		}
299
300		out = append(out, vm)
301		return out, nil
302	}
303
304	// List virtual machines
305	if len(args) == 0 {
306		return nil, errors.New("no argument")
307	}
308
309	finder, err := flag.Finder()
310	if err != nil {
311		return nil, err
312	}
313
314	var nfe error
315
316	// List virtual machines for every argument
317	for _, arg := range args {
318		vms, err := finder.VirtualMachineList(ctx, arg)
319		if err != nil {
320			if _, ok := err.(*find.NotFoundError); ok {
321				// Let caller decide how to handle NotFoundError
322				nfe = err
323				continue
324			}
325			return nil, err
326		}
327
328		out = append(out, vms...)
329	}
330
331	return out, nfe
332}
333
334func (flag *SearchFlag) VirtualApp() (*object.VirtualApp, error) {
335	ref, err := flag.search()
336	if err != nil {
337		return nil, err
338	}
339
340	app, ok := ref.(*object.VirtualApp)
341	if !ok {
342		return nil, fmt.Errorf("expected VirtualApp entity, got %s", ref.Reference().Type)
343	}
344
345	return app, nil
346}
347
348func (flag *SearchFlag) VirtualApps(args []string) ([]*object.VirtualApp, error) {
349	ctx := context.TODO()
350	var out []*object.VirtualApp
351
352	if flag.IsSet() {
353		app, err := flag.VirtualApp()
354		if err != nil {
355			return nil, err
356		}
357
358		out = append(out, app)
359		return out, nil
360	}
361
362	// List virtual apps
363	if len(args) == 0 {
364		return nil, errors.New("no argument")
365	}
366
367	finder, err := flag.Finder()
368	if err != nil {
369		return nil, err
370	}
371
372	// List virtual apps for every argument
373	for _, arg := range args {
374		apps, err := finder.VirtualAppList(ctx, arg)
375		if err != nil {
376			return nil, err
377		}
378
379		out = append(out, apps...)
380	}
381
382	return out, nil
383}
384
385func (flag *SearchFlag) HostSystem() (*object.HostSystem, error) {
386	ref, err := flag.search()
387	if err != nil {
388		return nil, err
389	}
390
391	host, ok := ref.(*object.HostSystem)
392	if !ok {
393		return nil, fmt.Errorf("expected HostSystem entity, got %s", ref.Reference().Type)
394	}
395
396	return host, nil
397}
398
399func (flag *SearchFlag) HostSystems(args []string) ([]*object.HostSystem, error) {
400	ctx := context.TODO()
401	var out []*object.HostSystem
402
403	if flag.IsSet() {
404		host, err := flag.HostSystem()
405		if err != nil {
406			return nil, err
407		}
408
409		out = append(out, host)
410		return out, nil
411	}
412
413	// List host system
414	if len(args) == 0 {
415		return nil, errors.New("no argument")
416	}
417
418	finder, err := flag.Finder()
419	if err != nil {
420		return nil, err
421	}
422
423	// List host systems for every argument
424	for _, arg := range args {
425		vms, err := finder.HostSystemList(ctx, arg)
426		if err != nil {
427			return nil, err
428		}
429
430		out = append(out, vms...)
431	}
432
433	return out, nil
434}
435