1/*
2Copyright (c) 2014-2017 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 guest
18
19import (
20	"context"
21	"flag"
22	"fmt"
23	"io"
24	"strconv"
25	"strings"
26	"text/tabwriter"
27	"time"
28
29	"github.com/vmware/govmomi/govc/cli"
30	"github.com/vmware/govmomi/govc/flags"
31	"github.com/vmware/govmomi/vim25/types"
32)
33
34type ps struct {
35	*flags.OutputFlag
36	*GuestFlag
37
38	every bool
39	exit  bool
40	wait  bool
41
42	pids pidSelector
43	uids uidSelector
44}
45
46type pidSelector []int64
47
48func (s *pidSelector) String() string {
49	return fmt.Sprint(*s)
50}
51
52func (s *pidSelector) Set(value string) error {
53	v, err := strconv.ParseInt(value, 0, 64)
54	if err != nil {
55		return err
56	}
57	*s = append(*s, v)
58	return nil
59}
60
61type uidSelector map[string]bool
62
63func (s uidSelector) String() string {
64	return ""
65}
66
67func (s uidSelector) Set(value string) error {
68	s[value] = true
69	return nil
70}
71
72func init() {
73	cli.Register("guest.ps", &ps{})
74}
75
76func (cmd *ps) Register(ctx context.Context, f *flag.FlagSet) {
77	cmd.OutputFlag, ctx = flags.NewOutputFlag(ctx)
78	cmd.OutputFlag.Register(ctx, f)
79
80	cmd.GuestFlag, ctx = newGuestFlag(ctx)
81	cmd.GuestFlag.Register(ctx, f)
82
83	cmd.uids = make(map[string]bool)
84	f.BoolVar(&cmd.every, "e", false, "Select all processes")
85	f.BoolVar(&cmd.exit, "x", false, "Output exit time and code")
86	f.BoolVar(&cmd.wait, "X", false, "Wait for process to exit")
87	f.Var(&cmd.pids, "p", "Select by process ID")
88	f.Var(&cmd.uids, "U", "Select by process UID")
89}
90
91func (cmd *ps) Process(ctx context.Context) error {
92	if err := cmd.OutputFlag.Process(ctx); err != nil {
93		return err
94	}
95	if err := cmd.GuestFlag.Process(ctx); err != nil {
96		return err
97	}
98	return nil
99}
100
101func (cmd *ps) Description() string {
102	return `List processes in VM.
103
104By default, unless the '-e', '-p' or '-U' flag is specified, only processes owned
105by the '-l' flag user are displayed.
106
107The '-x' and '-X' flags only apply to processes started by vmware-tools,
108such as those started with the govc guest.start command.
109
110Examples:
111  govc guest.ps -vm $name
112  govc guest.ps -vm $name -e
113  govc guest.ps -vm $name -p 12345
114  govc guest.ps -vm $name -U root`
115}
116
117func running(procs []types.GuestProcessInfo) bool {
118	for _, p := range procs {
119		if p.EndTime == nil {
120			return true
121		}
122	}
123	return false
124}
125
126func (cmd *ps) list(ctx context.Context) ([]types.GuestProcessInfo, error) {
127	m, err := cmd.ProcessManager()
128	if err != nil {
129		return nil, err
130	}
131
132	auth := cmd.Auth()
133
134	for {
135		procs, err := m.ListProcesses(ctx, auth, cmd.pids)
136		if err != nil {
137			return nil, err
138		}
139
140		if cmd.wait && running(procs) {
141			<-time.After(time.Millisecond * 250)
142			continue
143		}
144
145		return procs, nil
146	}
147}
148
149func (cmd *ps) Run(ctx context.Context, f *flag.FlagSet) error {
150	procs, err := cmd.list(ctx)
151	if err != nil {
152		return err
153	}
154
155	r := &psResult{cmd, procs}
156
157	return cmd.WriteResult(r)
158}
159
160type psResult struct {
161	cmd         *ps
162	ProcessInfo []types.GuestProcessInfo
163}
164
165func (r *psResult) Write(w io.Writer) error {
166	tw := tabwriter.NewWriter(w, 2, 0, 2, ' ', 0)
167
168	fmt.Fprintf(tw, "%s\t%s\t%s", "UID", "PID", "STIME")
169	if r.cmd.exit {
170		fmt.Fprintf(tw, "\t%s\t%s", "XTIME", "XCODE")
171	}
172	fmt.Fprint(tw, "\tCMD\n")
173
174	if len(r.cmd.pids) != 0 {
175		r.cmd.every = true
176	}
177
178	if !r.cmd.every && len(r.cmd.uids) == 0 {
179		r.cmd.uids[r.cmd.auth.Username] = true
180	}
181
182	for _, p := range r.ProcessInfo {
183		if r.cmd.every || r.cmd.uids[p.Owner] {
184			fmt.Fprintf(tw, "%s\t%d\t%s", p.Owner, p.Pid, p.StartTime.Format("15:04"))
185			if r.cmd.exit {
186				etime := "-"
187				ecode := "-"
188				if p.EndTime != nil {
189					etime = p.EndTime.Format("15:04")
190					ecode = strconv.Itoa(int(p.ExitCode))
191				}
192				fmt.Fprintf(tw, "\t%s\t%s", etime, ecode)
193			}
194			fmt.Fprintf(tw, "\t%s\n", strings.TrimSpace(p.CmdLine))
195		}
196	}
197
198	return tw.Flush()
199}
200