1 // DT PS Tree
2 //
3 // Douglas Thrift
4 //
5 // dtpstree.cpp
6 
7 /*  Copyright 2010 Douglas Thrift
8  *
9  *  Licensed under the Apache License, Version 2.0 (the "License");
10  *  you may not use this file except in compliance with the License.
11  *  You may obtain a copy of the License at
12  *
13  *      http://www.apache.org/licenses/LICENSE-2.0
14  *
15  *  Unless required by applicable law or agreed to in writing, software
16  *  distributed under the License is distributed on an "AS IS" BASIS,
17  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18  *  See the License for the specific language governing permissions and
19  *  limitations under the License.
20  */
21 
22 #include <cerrno>
23 #include <climits>
24 #include <clocale>
25 #include <cstdarg>
26 #include <cstdio>
27 #include <cstdlib>
28 #include <cstring>
29 #include <iostream>
30 #include <map>
31 #include <sstream>
32 #include <string>
33 #include <vector>
34 
35 #ifndef __GLIBC__
36 #include <libgen.h>
37 #endif
38 
39 #ifdef HAVE_TERMCAP_H
40 #include <termcap.h>
41 #elif defined(HAVE_NCURSES_TERMCAP_H)
42 #include <ncurses/termcap.h>
43 #elif defined(HAVE_NCURSES_TERM_H)
44 #include <ncurses/ncurses.h>
45 #include <ncurses/term.h>
46 #elif defined(HAVE_TERM_H)
47 #include <curses.h>
48 #include <term.h>
49 #endif
50 
51 #include <err.h>
52 #include <fcntl.h>
53 #include <getopt.h>
54 #include <kvm.h>
55 #include <paths.h>
56 #include <pwd.h>
57 #include <sys/param.h>
58 #include <sys/sysctl.h>
59 #include <sys/user.h>
60 #include <sys/utsname.h>
61 #include <unistd.h>
62 #include <vis.h>
63 
64 #include "foreach.hpp"
65 
66 namespace kvm
67 {
68 
69 #if HAVE_DECL_KERN_PROC_PROC
70 const int All(KERN_PROC_PROC);
71 #elif HAVE_DECL_KERN_PROC_KTHREAD
72 const int All(KERN_PROC_KTHREAD);
73 #else
74 const int All(KERN_PROC_ALL);
75 #endif
76 
77 template <typename Type>
78 inline Type *getprocs(kvm_t *kd, int &count);
79 
80 template <typename Type>
81 inline char **getargv(kvm_t *kd, const Type *proc);
82 
83 template <typename Type>
84 inline pid_t pid(Type *proc);
85 
86 template <typename Type>
87 inline pid_t ppid(Type *proc);
88 
89 template <typename Type>
90 inline uid_t ruid(Type *proc);
91 
92 template <typename Type>
93 inline char *comm(Type *proc);
94 
95 #ifndef HAVE_STRUCT_KINFO_PROC2
96 typedef kinfo_proc Proc;
97 
98 const int Flags(O_RDONLY);
99 
100 template <>
getprocs(kvm_t * kd,int & count)101 inline kinfo_proc *getprocs(kvm_t *kd, int &count)
102 {
103 	return kvm_getprocs(kd, All, 0, &count);
104 }
105 
106 template <>
getargv(kvm_t * kd,const kinfo_proc * proc)107 inline char **getargv(kvm_t *kd, const kinfo_proc *proc)
108 {
109 	return kvm_getargv(kd, proc, 0);
110 }
111 #else
112 typedef kinfo_proc2 Proc;
113 
114 const int Flags(KVM_NO_FILES);
115 
116 template <>
getprocs(kvm_t * kd,int & count)117 inline kinfo_proc2 *getprocs(kvm_t *kd, int &count)
118 {
119 	return kvm_getproc2(kd, All, 0, sizeof (kinfo_proc2), &count);
120 }
121 
122 template <>
getargv(kvm_t * kd,const kinfo_proc2 * proc)123 inline char **getargv(kvm_t *kd, const kinfo_proc2 *proc)
124 {
125 	return kvm_getargv2(kd, proc, 0);
126 }
127 #endif
128 
129 template <>
pid(Proc * proc)130 inline pid_t pid(Proc *proc)
131 {
132 #	ifdef HAVE_STRUCT_KINFO_PROCX_KI_PID
133 	return proc->ki_pid;
134 #	elif defined(HAVE_STRUCT_KINFO_PROCX_KP_PID)
135 	return proc->kp_pid;
136 #	elif defined(HAVE_STRUCT_KINFO_PROCX_P_PID)
137 	return proc->p_pid;
138 #	endif
139 }
140 
141 template <>
ppid(Proc * proc)142 inline pid_t ppid(Proc *proc)
143 {
144 #	ifdef HAVE_STRUCT_KINFO_PROCX_KI_PPID
145 	return proc->ki_ppid;
146 #	elif defined(HAVE_STRUCT_KINFO_PROCX_KP_PPID)
147 	return proc->kp_ppid;
148 #	elif defined(HAVE_STRUCT_KINFO_PROCX_P_PPID)
149 	return proc->p_ppid;
150 #	endif
151 }
152 
153 template <>
ruid(Proc * proc)154 inline uid_t ruid(Proc *proc)
155 {
156 #	ifdef HAVE_STRUCT_KINFO_PROCX_KI_RUID
157 	return proc->ki_ruid;
158 #	elif defined(HAVE_STRUCT_KINFO_PROCX_KP_RUID)
159 	return proc->kp_ruid;
160 #	elif defined(HAVE_STRUCT_KINFO_PROCX_P_RUID)
161 	return proc->p_ruid;
162 #	endif
163 }
164 
165 template <>
comm(Proc * proc)166 inline char *comm(Proc *proc)
167 {
168 #	ifdef HAVE_STRUCT_KINFO_PROCX_KI_COMM
169 	return proc->ki_comm;
170 #	elif defined(HAVE_STRUCT_KINFO_PROCX_KP_COMM)
171 	return proc->kp_comm;
172 #	elif defined(HAVE_STRUCT_KINFO_PROCX_P_COMM)
173 	return proc->p_comm;
174 #	endif
175 }
176 
177 }
178 
179 enum Flags
180 {
181 	Arguments	= 0x0001,
182 	Ascii		= 0x0002,
183 	NoCompact	= 0x0004,
184 	Highlight	= 0x0008,
185 	Vt100		= 0x0010,
186 	ShowKernel	= 0x0020,
187 	Long		= 0x0040,
188 	NumericSort	= 0x0080,
189 	ShowPids	= 0x0100,
190 	ShowTitles	= 0x0200,
191 	UidChanges	= 0x0400,
192 	Unicode		= 0x0800,
193 	Pid			= 0x1000,
194 	User		= 0x2000
195 };
196 
197 enum Escape { None, BoxDrawing, Bright };
198 
199 struct Segment
200 {
201 	size_t width_;
202 	Escape escape_;
203 	char *string_;
204 
SegmentSegment205 	inline Segment(size_t width, Escape escape, char *string) : width_(width), escape_(escape), string_(string) {}
206 };
207 
208 struct Branch
209 {
210 	std::string indentation_;
211 	bool done_;
212 
BranchBranch213 	inline Branch(size_t indentation) : indentation_(indentation, ' '), done_(false) {}
214 };
215 
216 class Tree
217 {
218 	const uint16_t &flags_;
219 	bool vt100_;
220 	wchar_t horizontal_, vertical_, upAndRight_, verticalAndRight_, downAndHorizontal_;
221 	size_t maxWidth_, width_;
222 	bool max_, suppress_;
223 	std::vector<Segment> segments_;
224 	std::vector<Branch> branches_;
225 	bool first_, last_;
226 	size_t duplicate_;
227 
228 public:
Tree(const uint16_t & flags)229 	Tree(const uint16_t &flags) : flags_(flags), vt100_(false), maxWidth_(0), width_(0), max_(false), suppress_(false), duplicate_(0)
230 	{
231 		bool tty(isatty(1));
232 
233 		if (flags & Ascii)
234 		{
235 		ascii:
236 			horizontal_ = L'-';
237 			vertical_ = L'|';
238 			upAndRight_ = L'`';
239 			verticalAndRight_ = L'|';
240 			downAndHorizontal_ = L'+';
241 		}
242 		else if (flags & Unicode)
243 		{
244 		unicode:
245 			if (!std::setlocale(LC_CTYPE, ""))
246 				goto vt100;
247 
248 			horizontal_ = L'\x2500';
249 			vertical_ = L'\x2502';
250 			upAndRight_ = L'\x2514';
251 			verticalAndRight_ = L'\x251c';
252 			downAndHorizontal_ = L'\x252c';
253 
254 			wchar_t wides[] = { horizontal_, vertical_, upAndRight_, verticalAndRight_, downAndHorizontal_ };
255 			char *buffer = new char[MB_CUR_MAX];
256 
257 			for (int index(0); index != sizeof (wides) / sizeof (*wides); ++index)
258 			{
259 				int size;
260 
261 				if ((size = std::wctomb(buffer, wides[index])) == -1)
262 				{
263 					delete [] buffer;
264 					goto vt100;
265 				}
266 
267 				wchar_t wide;
268 
269 				if (std::mbtowc(&wide, buffer, size) == -1)
270 				{
271 					delete [] buffer;
272 					goto vt100;
273 				}
274 
275 				if (wide != wides[index])
276 				{
277 					delete [] buffer;
278 					goto vt100;
279 				}
280 			}
281 
282 			delete [] buffer;
283 		}
284 		else if (flags & Vt100)
285 		{
286 		vt100:
287 			vt100_ = true;
288 			horizontal_ = L'\x71';
289 			vertical_ = L'\x78';
290 			upAndRight_ = L'\x6d';
291 			verticalAndRight_ = L'\x74';
292 			downAndHorizontal_ = L'\x77';
293 		}
294 		else if (tty)
295 			goto unicode;
296 		else
297 			goto ascii;
298 
299 		if (!(flags & Long) && tty)
300 		{
301 #			if !defined(HAVE_TERMCAP_H) && !defined(HAVE_NCURSES_TERMCAP_H)
302 			int code;
303 
304 			if (setupterm(NULL, 1, &code) == OK)
305 			{
306 				maxWidth_ = tigetnum(const_cast<char *>("cols"));
307 
308 				if (tigetflag(const_cast<char *>("am")) && !tigetflag(const_cast<char *>("xenl")))
309 					suppress_ = true;
310 			}
311 #			else
312 			char buffer[1024], *term(std::getenv("TERM"));
313 
314 			if (term != NULL && tgetent(buffer, term) == 1)
315 			{
316 				maxWidth_ = tgetnum("co");
317 
318 				if (tgetflag("am") && !tgetflag("xn"))
319 					suppress_ = true;
320 			}
321 #			endif
322 			else
323 				maxWidth_ = 80;
324 		}
325 	}
326 
print(const std::string & string,bool highlight,size_t duplicate)327 	void print(const std::string &string, bool highlight, size_t duplicate)
328 	{
329 		Escape escape(vt100_ ? BoxDrawing : None);
330 
331 		if (!first_ || flags_ & Arguments)
332 		{
333 			size_t last(branches_.size() - 1);
334 
335 			_foreach (std::vector<Branch>, branch, branches_)
336 			{
337 				size_t width(branch->indentation_.size() + 2);
338 
339 				if (_index == last)
340 				{
341 					wchar_t line;
342 
343 					if (last_)
344 					{
345 						branch->done_ = true;
346 						line = upAndRight_;
347 					}
348 					else
349 						line = verticalAndRight_;
350 
351 					print(width, escape, "%s%lc%lc", branch->indentation_.c_str(), line, horizontal_);
352 				}
353 				else
354 					print(width, escape, "%s%lc ", branch->indentation_.c_str(), branch->done_ ? ' ' : vertical_);
355 			}
356 		}
357 		else if (branches_.size())
358 		{
359 			wchar_t line;
360 
361 			if (last_)
362 			{
363 				branches_.back().done_ = true;
364 				line = horizontal_;
365 			}
366 			else
367 				line = downAndHorizontal_;
368 
369 			print(3, escape, "%lc%lc%lc", horizontal_, line, horizontal_);
370 		}
371 
372 		size_t size(0);
373 
374 		if (duplicate)
375 		{
376 			std::ostringstream string;
377 
378 			string << duplicate << "*[";
379 
380 			size = string.str().size();
381 
382 			print(size, None, "%s", string.str().c_str());
383 
384 			++duplicate_;
385 		}
386 
387 		print(string.size(), highlight ? Bright : None, "%s", string.c_str());
388 
389 		branches_.push_back(Branch(!(flags_ & Arguments) ? size + string.size() + 1 : 2));
390 	}
391 
printArg(const std::string & arg,bool last)392 	inline void printArg(const std::string &arg, bool last)
393 	{
394 		if (max_)
395 			return;
396 
397 		size_t width(arg.size() + 1);
398 
399 		width_ += width;
400 
401 		char *string;
402 
403 		if (maxWidth_ && !(flags_ & Long))
404 			if (width_ > maxWidth_ || !last && width_ + 3 >= maxWidth_)
405 			{
406 				width -= width_ - maxWidth_;
407 				width_ = maxWidth_;
408 				max_ = true;
409 
410 				ssize_t size(static_cast<ssize_t>(width) - 4);
411 
412 				if (size < -3)
413 					return;
414 				else if (size < 1)
415 				{
416 					string = static_cast<char *>(std::malloc(size += 5));
417 
418 					snprintf(string, size, " ...");
419 				}
420 				else
421 					asprintf(&string, " %s...", arg.substr(0, size).c_str());
422 			}
423 			else
424 				goto print;
425 		else
426 		print:
427 			asprintf(&string, " %s", arg.c_str());
428 
429 		segments_.push_back(Segment(width, None, string));
430 	}
431 
pop(bool children)432 	inline void pop(bool children)
433 	{
434 		branches_.pop_back();
435 
436 		if (!(flags_ & Arguments) && !children)
437 			done();
438 	}
439 
done()440 	void done()
441 	{
442 		if (duplicate_)
443 		{
444 			print(duplicate_, None, "%s", std::string(duplicate_, ']').c_str());
445 
446 			duplicate_ = 0;
447 		}
448 
449 		size_t last(segments_.size() - 1);
450 
451 		_foreach (std::vector<Segment>, segment, segments_)
452 		{
453 			const char *begin, *end;
454 
455 			switch (segment->escape_)
456 			{
457 			case BoxDrawing:
458 				begin = !_index || (segment - 1)->escape_ != BoxDrawing ? "\033(0\017" : "";
459 				end = _index == last || (segment + 1)->escape_ != BoxDrawing ? "\033(B\017" : "";
460 				break;
461 			case Bright:
462 				begin = "\033[1m";
463 				end = "\033[22m";
464 				break;
465 			default:
466 				begin = end = ""; break;
467 			}
468 
469 			std::printf("%s%s%s", begin, segment->string_, end);
470 			std::free(segment->string_);
471 		}
472 
473 		segments_.clear();
474 
475 		if (suppress_ && width_ == maxWidth_)
476 			std::fflush(stdout);
477 		else
478 			std::printf("\n");
479 
480 		width_ = 0;
481 		max_ = false;
482 	}
483 
operator ()(bool first,bool last)484 	inline Tree &operator()(bool first, bool last)
485 	{
486 		first_ = first;
487 		last_ = last;
488 
489 		return *this;
490 	}
491 
492 private:
print(size_t width,Escape escape,const char * format,...)493 	void print(size_t width, Escape escape, const char * format, ...)
494 	{
495 		if (max_)
496 			return;
497 
498 		std::va_list args;
499 
500 		va_start(args, format);
501 
502 		char *string;
503 
504 		vasprintf(&string, format, args);
505 		va_end(args);
506 
507 		width_ += width;
508 
509 		if (maxWidth_ && !(flags_ & Long))
510 			if (width_ > maxWidth_)
511 			{
512 				width -= width_ - maxWidth_;
513 				width_ = maxWidth_;
514 				max_ = true;
515 
516 				bool previous = !width;
517 
518 				if (previous)
519 				{
520 					std::free(string);
521 
522 					const Segment &segment(segments_.back());
523 
524 					width = segment.width_;
525 					string = segment.string_;
526 				}
527 
528 				std::wstring wide(width - 1, '\0');
529 
530 				std::mbstowcs(const_cast<wchar_t *>(wide.data()), string, wide.size());
531 				std::free(string);
532 
533 				wide += L'+';
534 
535 				size_t size(std::wcstombs(NULL, wide.c_str(), 0) + 1);
536 
537 				string = static_cast<char *>(std::malloc(size));
538 
539 				std::wcstombs(string, wide.c_str(), size);
540 
541 				if (previous)
542 				{
543 					segments_.back().string_ = string;
544 
545 					return;
546 				}
547 			}
548 
549 		segments_.push_back(Segment(width, escape, string));
550 	}
551 };
552 
553 template <typename Type>
554 struct Proc
555 {
556 	typedef std::multimap<pid_t, Proc<Type> *> PidMap;
557 	typedef std::multimap<std::string, Proc<Type> *> NameMap;
558 
559 private:
560 	const uint16_t &flags_;
561 	kvm_t *kd_;
562 	Type *proc_;
563 	mutable std::string name_, print_;
564 	Proc<Type> *parent_;
565 	PidMap childrenByPid_;
566 	NameMap childrenByName_;
567 	bool highlight_, root_;
568 	int8_t compact_;
569 	size_t duplicate_;
570 
571 public:
ProcProc572 	inline Proc(const uint16_t &flags, kvm_t *kd, Type *proc) : flags_(flags), kd_(kd), proc_(proc), parent_(NULL), highlight_(false), root_(false), compact_(-1), duplicate_(0) {}
573 
nameProc574 	inline const std::string &name() const
575 	{
576 		if (name_.empty())
577 			name_ = visual(kvm::comm(proc_));
578 
579 		return name_;
580 	}
581 
parentProc582 	inline pid_t parent() const { return kvm::ppid(proc_); }
pidProc583 	inline pid_t pid() const { return kvm::pid(proc_); }
584 
childProc585 	inline void child(Proc *proc)
586 	{
587 		if (proc == this)
588 			return;
589 
590 		proc->parent_ = this;
591 
592 		childrenByPid_.insert(typename PidMap::value_type(proc->pid(), proc));
593 		childrenByName_.insert(typename NameMap::value_type(proc->name(), proc));
594 	}
595 
highlightProc596 	inline void highlight()
597 	{
598 		highlight_ = true;
599 
600 		if (parent_)
601 			parent_->highlight();
602 	}
603 
compactProc604 	inline bool compact()
605 	{
606 		if (compact_ == -1)
607 			compact_ = compact(childrenByName_);
608 
609 		return compact_;
610 	}
611 
rootProc612 	bool root(uid_t uid)
613 	{
614 		if (flags_ & User)
615 		{
616 			if (uid == this->uid())
617 			{
618 				Proc *parent(parent_);
619 
620 				while (parent)
621 				{
622 					if (parent->uid() == uid)
623 						return false;
624 
625 					parent = parent->parent_;
626 				}
627 
628 				return root_ = true;
629 			}
630 
631 			return false;
632 		}
633 
634 		return root_ = !parent_;
635 	}
636 
printByPidProc637 	inline void printByPid(Tree &tree) const
638 	{
639 		print(tree, childrenByPid_);
640 	}
641 
printByNameProc642 	inline void printByName(Tree &tree) const
643 	{
644 		print(tree, childrenByName_);
645 	}
646 
compactProc647 	static bool compact(NameMap &names)
648 	{
649 		Proc *previous(NULL);
650 		bool compact(true);
651 
652 		_tforeach (NameMap, name, names)
653 		{
654 			Proc *proc(name->second);
655 
656 			if (proc->duplicate_)
657 				continue;
658 
659 			size_t duplicate(proc->compact());
660 
661 			if (compact && duplicate && (!previous || proc->print() == previous->print()))
662 				previous = proc;
663 			else
664 				compact = false;
665 
666 			size_t count(names.count(name->first));
667 
668 			if (!duplicate || count == 1)
669 				continue;
670 
671 			_forall(typename NameMap::iterator, n4me, (++name)--, names.upper_bound(name->first))
672 			{
673 				Proc *pr0c(n4me->second);
674 
675 				if (pr0c->compact() && Proc::compact(proc, pr0c))
676 					duplicate += ++pr0c->duplicate_;
677 			}
678 
679 			if (duplicate != 1)
680 				proc->duplicate_ = duplicate;
681 		}
682 
683 		return compact;
684 	}
685 
686 private:
visualProc687 	inline std::string visual(const char *string) const
688 	{
689 		std::string visual(std::strlen(string) * 4 + 1, '\0');
690 
691 		visual.resize(strvis(const_cast<char *>(visual.data()), string, VIS_TAB | VIS_NL | VIS_NOSLASH));
692 
693 		return visual;
694 	}
695 
696 	template <typename Map>
printProc697 	void print(Tree &tree, const Map &children) const
698 	{
699 		if (duplicate_ == 1)
700 			return;
701 
702 		print(tree);
703 
704 		size_t size(children.size()), last(size - 1);
705 
706 		_tforeach (const Map, child, children)
707 		{
708 			Proc<Type> *proc(child->second);
709 			bool l4st(_index + (proc->duplicate_ ? proc->duplicate_ - 1 : 0) == last);
710 
711 			if (!l4st)
712 			{
713 				l4st = true;
714 
715 				for (++child; child != _end; ++child)
716 					if (child->second->duplicate_ != 1)
717 					{
718 						l4st = false;
719 
720 						break;
721 					}
722 
723 				--child;
724 			}
725 
726 			proc->print(tree(!_index, l4st), proc->template children<Map>());
727 
728 			if (l4st)
729 				break;
730 		}
731 
732 		tree.pop(size);
733 	}
734 
printProc735 	void print(Tree &tree) const
736 	{
737 		tree.print(print(), highlight_, duplicate_);
738 
739 		if (flags_ & Arguments)
740 		{
741 			char **argv(kvm::getargv(kd_, proc_));
742 
743 			if (argv && *argv)
744 				for (++argv; *argv; ++argv)
745 					tree.printArg(visual(*argv), !*(argv + 1));
746 
747 			tree.done();
748 		}
749 	}
750 
printProc751 	const std::string &print() const
752 	{
753 		if (print_.empty())
754 		{
755 			std::ostringstream print;
756 
757 			if (flags_ & ShowTitles)
758 			{
759 				char **argv(kvm::getargv(kd_, proc_));
760 
761 				if (argv)
762 					print << visual(*argv);
763 				else
764 					print << name();
765 			}
766 			else
767 				print << name();
768 
769 			bool p1d(flags_ & ShowPids), args(flags_ & Arguments);
770 			bool change(flags_ & UidChanges && (root_ ? !(flags_ & User) && uid() : parent_ && uid() != parent_->uid()));
771 			bool parens((p1d || change) && !args);
772 
773 			if (parens)
774 				print << '(';
775 
776 			if (p1d)
777 			{
778 				if (!parens)
779 					print << ',';
780 
781 				print << pid();
782 			}
783 
784 			if (change)
785 			{
786 				if (!parens || p1d)
787 					print << ',';
788 
789 				passwd *user(getpwuid(uid()));
790 
791 				print << user->pw_name;
792 			}
793 
794 			if (parens)
795 				print << ')';
796 
797 			print_ = print.str();
798 		}
799 
800 		return print_;
801 	}
802 
uidProc803 	inline uid_t uid() const { return kvm::ruid(proc_); }
804 
805 	template <typename Map>
806 	inline const Map &children() const;
807 
hasChildrenProc808 	inline bool hasChildren() const { return childrenByName_.size(); }
childProc809 	inline Proc<Type> *child() const { return childrenByName_.begin()->second; }
810 
compactProc811 	inline static bool compact(Proc<Type> *one, Proc<Type> *two)
812 	{
813 		if (one->print() != two->print())
814 			return false;
815 
816 		if (one->hasChildren() != two->hasChildren())
817 			return false;
818 
819 		if (one->hasChildren() && !compact(one->child(), two->child()))
820 			return false;
821 
822 		if (two->highlight_)
823 			one->highlight_ = true;
824 
825 		return true;
826 	}
827 };
828 
829 template <> template <>
children() const830 inline const Proc<kvm::Proc>::PidMap &Proc<kvm::Proc>::children() const
831 {
832 	return childrenByPid_;
833 }
834 
835 template <> template <>
children() const836 inline const Proc<kvm::Proc>::NameMap &Proc<kvm::Proc>::children() const
837 {
838 	return childrenByName_;
839 }
840 
help(char * program,option options[],int code=0)841 static void help(char *program, option options[], int code = 0)
842 {
843 	std::printf("Usage: %s [options] [PID|USER]\n\nOptions:\n", basename(program));
844 
845 	for (option *option(options); option->name; ++option)
846 	{
847 		std::string name(option->name);
848 		std::ostringstream arguments;
849 
850 		switch (option->val)
851 		{
852 		case 'H':
853 			if (name != "highlight")
854 				continue;
855 
856 			arguments << "-H[PID], --highlight[=PID]"; break;
857 		case 0:
858 			if (name == "pid")
859 				arguments << "PID, --pid=PID";
860 			else if (name == "user")
861 				arguments << "USER, --user=USER";
862 			else
863 				goto argument;
864 
865 			break;
866 		case 'c':
867 			if (name != "no-compact")
868 				continue;
869 		default:
870 			arguments << '-' << static_cast<char>(option->val) << ", ";
871 		argument:
872 			arguments << "--" << name;
873 		}
874 
875 		const char *description("");
876 
877 		switch (option->val)
878 		{
879 		case 'a':
880 			description = "show command line arguments"; break;
881 		case 'A':
882 			description = "use ASCII line drawing characters"; break;
883 		case 'c':
884 			description = "don't compact identical subtrees"; break;
885 		case 'h':
886 			description = "show this help message and exit"; break;
887 		case 'H':
888 			description = "highlight the current process (or PID) and its\n                              ancestors"; break;
889 		case 'G':
890 			description = "use VT100 line drawing characters"; break;
891 		case 'k':
892 			description = "show kernel processes"; break;
893 		case 'l':
894 			description = "don't truncate long lines"; break;
895 		case 'n':
896 			description = "sort output by PID"; break;
897 		case 'p':
898 			description = "show PIDs; implies -c"; break;
899 		case 't':
900 			description = "show process titles"; break;
901 		case 'u':
902 			description = "show uid transitions"; break;
903 		case 'U':
904 			description = "use Unicode line drawing characters"; break;
905 		case 'V':
906 			description = "show version information and exit"; break;
907 		case 0:
908 			if (name == "pid")
909 				description = "show only the tree rooted at the process PID";
910 			else if (name == "user")
911 				description = "show only trees rooted at processes of USER";
912 		}
913 
914 		std::printf("  %-27s %s\n", arguments.str().c_str(), description);
915 	}
916 
917 	std::exit(code);
918 }
919 
920 template <typename Type, long minimum, long maximum>
value(char * program,option options[],bool * success=NULL)921 static Type value(char *program, option options[], bool *success = NULL)
922 {
923 	char *end;
924 	long value(std::strtol(optarg, &end, 0));
925 
926 	errno = 0;
927 
928 	if (optarg == end || *end != '\0')
929 		if (success)
930 			*success = false;
931 		else
932 		{
933 			warnx("Number is invalid: \"%s\"", optarg);
934 			help(program, options, 1);
935 		}
936 	else if (value < minimum || value == LONG_MIN && errno == ERANGE)
937 	{
938 		warnx("Number is too small: \"%s\"", optarg);
939 		help(program, options, 1);
940 	}
941 	else if (value > maximum || value == LONG_MAX && errno == ERANGE)
942 	{
943 		warnx("Number is too large: \"%s\"", optarg);
944 		help(program, options, 1);
945 	}
946 	else if (success)
947 		*success = true;
948 
949 	return value;
950 }
951 
options(int argc,char * argv[],pid_t & hpid,pid_t & pid,char * & user)952 static uint16_t options(int argc, char *argv[], pid_t &hpid, pid_t &pid, char *&user)
953 {
954 	option options[] = {
955 		{ "arguments", no_argument, NULL, 'a' },
956 		{ "ascii", no_argument, NULL, 'A' },
957 		{ "compact", no_argument, NULL, 'c' },
958 		{ "no-compact", no_argument, NULL, 'c' },
959 		{ "help", no_argument, NULL, 'h' },
960 		{ "highlight", optional_argument, NULL, 'H' },
961 		{ "highlight-all", no_argument, NULL, 'H' },
962 		{ "highlight-pid", required_argument, NULL, 'H' },
963 		{ "vt100", no_argument, NULL, 'G' },
964 		{ "show-kernel", no_argument, NULL, 'k' },
965 		{ "long", no_argument, NULL, 'l' },
966 		{ "numeric-sort", no_argument, NULL, 'n' },
967 		{ "show-pids", no_argument, NULL, 'p' },
968 		{ "show-titles", no_argument, NULL, 't' },
969 		{ "uid-changes", no_argument, NULL, 'u' },
970 		{ "unicode", no_argument, NULL, 'U' },
971 		{ "version", optional_argument, NULL, 'V' },
972 		{ "pid", required_argument, NULL, 0 },
973 		{ "user", required_argument, NULL, 0 },
974 		{ NULL, 0, NULL, 0 }
975 	};
976 	int option, index;
977 	uint16_t flags(0);
978 	char *program(argv[0]);
979 
980 	while ((option = getopt_long(argc, argv, "aAchH::GklnptuUV::", options, &index)) != -1)
981 		switch (option)
982 		{
983 		case 'a':
984 			flags |= Arguments | NoCompact; break;
985 		case 'A':
986 			flags |= Ascii;
987 			flags &= ~Vt100;
988 			flags &= ~Unicode;
989 
990 			break;
991 		case 'c':
992 			flags |= NoCompact; break;
993 		case 'h':
994 			help(program, options);
995 		case 'H':
996 			hpid = optarg ? value<pid_t, 0, INT_MAX>(program, options) : getpid();
997 			flags |= Highlight;
998 
999 			break;
1000 		case 'G':
1001 			flags |= Vt100;
1002 			flags &= ~Ascii;
1003 			flags &= ~Unicode;
1004 
1005 			break;
1006 		case 'k':
1007 			flags |= ShowKernel; break;
1008 		case 'l':
1009 			flags |= Long; break;
1010 		case 'n':
1011 			flags |= NumericSort; break;
1012 		case 'p':
1013 			flags |= NoCompact | ShowPids; break;
1014 		case 't':
1015 			flags |= ShowTitles; break;
1016 		case 'u':
1017 			flags |= UidChanges; break;
1018 		case 'U':
1019 			flags |= Unicode;
1020 			flags &= ~Ascii;
1021 			flags &= ~Vt100;
1022 
1023 			break;
1024 		case 'V':
1025 			{
1026 				std::string version(optarg ? optarg : "");
1027 
1028 				if (version == "s" || version == "short")
1029 					std::printf(PACKAGE_TARNAME " " PACKAGE_VERSION "\n");
1030 				else
1031 				{
1032 					utsname name;
1033 
1034 					if (uname(&name))
1035 						err(1, NULL);
1036 
1037 					std::printf(PACKAGE_TARNAME " " PACKAGE_VERSION " - %s %s %s\n", name.sysname, name.release, name.machine);
1038 				}
1039 
1040 				if (version == "l" || version == "license")
1041 					std::printf("\n"
1042 						"   Copyright 2010 Douglas Thrift\n\n"
1043 						"   Licensed under the Apache License, Version 2.0 (the \"License\");\n"
1044 						"   you may not use this file except in compliance with the License.\n"
1045 						"   You may obtain a copy of the License at\n\n"
1046 						"       http://www.apache.org/licenses/LICENSE-2.0\n\n"
1047 						"   Unless required by applicable law or agreed to in writing, software\n"
1048 						"   distributed under the License is distributed on an \"AS IS\" BASIS,\n"
1049 						"   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n"
1050 						"   See the License for the specific language governing permissions and\n"
1051 						"   limitations under the License.\n");
1052 
1053 				std::exit(0);
1054 			}
1055 		case 0:
1056 			{
1057 				std::string option(options[index].name);
1058 
1059 				if (option == "pid")
1060 				{
1061 					pid = value<pid_t, 0, INT_MAX>(program, options);
1062 					flags |= Pid;
1063 					flags &= ~User;
1064 				}
1065 				else if (option == "user")
1066 				{
1067 					std::free(user);
1068 
1069 					user = strdup(optarg);
1070 					flags |= User;
1071 					flags &= ~Pid;
1072 				}
1073 			}
1074 
1075 			break;
1076 		case '?':
1077 			help(program, options, 1);
1078 		}
1079 
1080 	_forall (int, index, optind, argc)
1081 	{
1082 		bool success(false);
1083 
1084 		optarg = argv[index];
1085 		pid = value<pid_t, 0, INT_MAX>(program, options, &success);
1086 
1087 		if (success)
1088 		{
1089 			flags |= Pid;
1090 			flags &= ~User;
1091 		}
1092 		else
1093 		{
1094 			std::free(user);
1095 
1096 			user = strdup(optarg);
1097 			flags |= User;
1098 			flags &= ~Pid;
1099 		}
1100 	}
1101 
1102 	return flags;
1103 }
1104 
1105 template <typename Type, int Flags>
tree(pid_t hpid,pid_t pid,uint16_t flags,uid_t uid)1106 static void tree(pid_t hpid, pid_t pid, uint16_t flags, uid_t uid)
1107 {
1108 	char error[_POSIX2_LINE_MAX];
1109 	kvm_t *kd(kvm_openfiles(NULL, _PATH_DEVNULL, NULL, Flags, error));
1110 
1111 	if (!kd)
1112 		errx(1, "%s", error);
1113 
1114 	int count;
1115 	Type *procs(kvm::getprocs<Type>(kd, count));
1116 
1117 	if (!procs)
1118 		errx(1, "%s", kvm_geterr(kd));
1119 
1120 	typedef Type *Pointer;
1121 	typename Proc<Type>::PidMap pids;
1122 
1123 	_forall (Pointer, proc, procs, procs + count)
1124 		if (flags & ShowKernel || kvm::ppid(proc) > 0 || kvm::pid(proc) == 1)
1125 			pids.insert(typename Proc<Type>::PidMap::value_type(kvm::pid(proc), new Proc<Type>(flags, kd, proc)));
1126 
1127 	enum { PidSort, NameSort } sort(flags & NumericSort ? PidSort : NameSort);
1128 
1129 	_tforeach (typename Proc<Type>::PidMap, pid, pids)
1130 	{
1131 		Proc<Type> *proc(pid->second);
1132 
1133 		if (proc->parent() == -1)
1134 			continue;
1135 
1136 		typename Proc<Type>::PidMap::iterator parent(pids.find(proc->parent()));
1137 
1138 		if (parent != pids.end())
1139 			parent->second->child(proc);
1140 	}
1141 
1142 	if (flags & Highlight)
1143 	{
1144 		typename Proc<Type>::PidMap::iterator pid(pids.find(hpid));
1145 
1146 		if (pid != pids.end())
1147 			pid->second->highlight();
1148 	}
1149 
1150 	Tree tree(flags);
1151 
1152 	if (flags & Pid)
1153 	{
1154 		typename Proc<Type>::PidMap::iterator p1d(pids.find(pid));
1155 
1156 		if (p1d != pids.end())
1157 		{
1158 			Proc<Type> *proc(p1d->second);
1159 
1160 			if (!(flags & NoCompact))
1161 				proc->compact();
1162 
1163 			switch (sort)
1164 			{
1165 			case PidSort:
1166 				proc->printByPid(tree);
1167 
1168 				break;
1169 			case NameSort:
1170 				proc->printByName(tree);
1171 			}
1172 		}
1173 	}
1174 	else
1175 	{
1176 		typename Proc<Type>::NameMap names;
1177 
1178 		_tforeach (typename Proc<Type>::PidMap, pid, pids)
1179 		{
1180 			Proc<Type> *proc(pid->second);
1181 
1182 			if (proc->root(uid))
1183 				names.insert(typename Proc<Type>::NameMap::value_type(proc->name(), proc));
1184 		}
1185 
1186 		if (!(flags & NoCompact))
1187 			Proc<Type>::compact(names);
1188 
1189 		switch (sort)
1190 		{
1191 		case PidSort:
1192 			_tforeach (typename Proc<Type>::PidMap, pid, pids)
1193 			{
1194 				Proc<Type> *proc(pid->second);
1195 
1196 				if (proc->root(uid))
1197 					proc->printByPid(tree);
1198 			}
1199 
1200 			break;
1201 		case NameSort:
1202 			_tforeach (typename Proc<Type>::NameMap, name, names)
1203 				name->second->printByName(tree);
1204 		}
1205 	}
1206 
1207 	_tforeach (typename Proc<Type>::PidMap, pid, pids)
1208 		delete pid->second;
1209 }
1210 
main(int argc,char * argv[])1211 int main(int argc, char *argv[])
1212 {
1213 	pid_t hpid(0), pid(0);
1214 	char *user(NULL);
1215 	uint16_t flags(options(argc, argv, hpid, pid, user));
1216 	uid_t uid(0);
1217 
1218 	if (flags & User)
1219 	{
1220 		errno = 0;
1221 
1222 		passwd *us3r(getpwnam(user));
1223 
1224 		if (!us3r)
1225 			errno ? err(1, NULL) : errx(1, "Unknown user: \"%s\"", user);
1226 
1227 		uid = us3r->pw_uid;
1228 	}
1229 
1230 	tree<kvm::Proc, kvm::Flags>(hpid, pid, flags, uid);
1231 
1232 	return 0;
1233 }
1234 
1235 // display a tree of processes
1236