Progress...

GNUmakefile

@@ -1,3 +1,4 @@
+CXXFLAGS := -pedantic -Wall
 LDFLAGS := -lkvm
 
 all: pstree

foreach.hpp

@@ -0,0 +1,106 @@
+// Foreach for STL
+//
+// Douglas Thrift
+//
+// $Id: foreach.hpp 1263 2010-03-21 13:30:15Z douglas $
+
+/* Menes - C++ High-Level Utility Library
+ * Copyright (C) 2004  Jay Freeman (saurik)
+*/
+
+/*
+ *        Redistribution and use in source and binary
+ * forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ * 
+ * 1. Redistributions of source code must retain the
+ *    above copyright notice, this list of conditions
+ *    and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the
+ *    above copyright notice, this list of conditions
+ *    and the following disclaimer in the documentation
+ *    and/or other materials provided with the
+ *    distribution.
+ * 3. The name of the author may not be used to endorse
+ *    or promote products derived from this software
+ *    without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS''
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING,
+ * BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE
+ * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+ * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
+ * TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+ * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
+ * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
+#ifndef _foreach_hpp_
+#define _foreach_hpp_
+
+#define _forever for (;;)
+
+#define _forall(type, item, begin, end) \
+	if (size_t _index = 0); \
+	else for (type item(begin), _end(end); item != _end; ++item, ++_index)
+
+#define _rforall(type, item, begin, end) \
+	for (type item(end), _begin(begin); item != _begin && (--item, true); )
+
+#define _repeat(count) \
+	for (unsigned _index(0), _end(count); _index != _end; ++_index)
+
+template <typename List_, bool noop = true>
+struct StrictIterator;
+
+template <typename List_>
+struct StrictIterator<List_, true> {
+	typedef typename List_::iterator Result;
+};
+
+template <typename List_>
+struct StrictIterator<const List_, true> {
+	typedef typename List_::const_iterator Result;
+};
+
+template <typename List_, bool noop = true>
+struct ListTraits;
+
+template <typename List_>
+struct ListTraits<List_, true>
+{
+	typedef typename StrictIterator<List_>::Result Iterator;
+
+	static inline Iterator Begin(List_ &arg) {
+		return arg.begin();
+	}
+
+	static inline Iterator End(List_ &arg) {
+		return arg.end();
+	}
+};
+
+#define _foreach_(type, item, set, forall, _typename) \
+	for (bool _stop(true); _stop; ) \
+		for (type &_set = set; _stop; _stop = false) \
+			forall (_typename ListTraits< type >::Iterator, item, ListTraits< type >::Begin(_set), ListTraits< type >::End(_set))
+
+#define _foreach(type, item, set) \
+	_foreach_(type, item, set, _forall, )
+
+#define _rforeach(type, item, set) \
+	_foreach_(type, item, set, _rforall, )
+
+#define _tforeach(type, item, set) \
+	_foreach_(type, item, set, _forall, typename)
+
+#define _rtforeach(type, item, set) \
+	_foreach_(type, item, set, _rforall, typename)
+
+#endif//_foreach_hpp_

pstree.cpp

@@ -4,32 +4,227 @@
 //
 // $Id$
 
+#include <climits>
+#include <cstdio>
+#include <iostream>
+#include <map>
+#include <set>
+#include <sstream>
+#include <string>
+
 #include <err.h>
 #include <fcntl.h>
 #include <kvm.h>
-#include <limits.h>
 #include <paths.h>
-#include <stdio.h>
+#include <pwd.h>
 #include <sys/param.h>
 #include <sys/sysctl.h>
 #include <sys/user.h>
 
+#include "foreach.hpp"
+
+class Proc;
+
+typedef std::map<pid_t, Proc *> PidMap;
+typedef std::multimap<std::string, Proc *> NameMap;
+
+enum Sort { Pid, Name };
+
+class Proc
+{
+	kvm_t *kd_;
+	kinfo_proc *proc_;
+	Proc *parent_;
+	PidMap childrenByPid_;
+	NameMap childrenByName_;
+	bool highlight_;
+
+public:
+	Proc() {}
+	Proc(kvm_t *kd, kinfo_proc *proc) : kd_(kd), proc_(proc) {}
+
+	inline std::string name() const { return proc_->ki_comm; }
+	inline pid_t parent() const { return proc_->ki_ppid; }
+	inline pid_t pid() const { return proc_->ki_pid; }
+
+	void child(Proc *proc)
+	{
+		proc->parent_ = this;
+		childrenByPid_[proc->pid()] = proc;
+
+		childrenByName_.insert(NameMap::value_type(proc->name(), proc));
+	}
+
+	void highlight()
+	{
+		highlight_ = true;
+
+		if (parent_)
+			parent_->highlight();
+	}
+
+	inline bool root() const { return !parent_; }
+
+	void printByPid(const std::string &indent = "") const
+	{
+		std::cout << indent << print(indent.size()) << std::endl;
+
+		_foreach (const PidMap, child, childrenByPid_)
+			child->second->printByPid(indent + "  ");
+	}
+
+	void printByName(const std::string &indent = "") const
+	{
+		std::cout << indent << print(indent.size()) << std::endl;
+
+		_foreach (const NameMap, child, childrenByName_)
+			child->second->printByName(indent + "  ");
+	}
+
+private:
+	std::string print(std::string::size_type indent) const
+	{
+		std::ostringstream print;
+
+		if (highlight_)
+			print << "\033[1m";
+
+		print << name();
+
+		bool _pid(true), _args(true);
+		bool change(true && parent_ && uid() != parent_->uid());
+		bool parens((_pid || change) && !_args);
+
+		if (parens)
+			print << '(';
+
+		if (_pid)
+		{
+			if (!parens)
+				print << ',';
+
+			print << pid();
+		}
+
+		if (change)
+		{
+			if (!parens || _pid)
+				print << ',';
+
+			print << user();
+		}
+
+		if (parens)
+			print << ')';
+
+		if (highlight_)
+			print << "\033[22m";
+
+		if (_args)
+			print << args(indent + print.str().size());
+
+		return print.str();
+	}
+
+	inline bool children() const { return childrenByPid_.size(); }
+
+	inline uid_t uid() const { return proc_->ki_ruid; }
+
+	std::string user() const
+	{
+		passwd *user(getpwuid(uid()));
+
+		return user->pw_name;
+	}
+
+	std::string args(std::string::size_type indent) const
+	{
+		char **argv(kvm_getargv(kd_, proc_, 0));
+		std::ostringstream args;
+
+		if (argv && *argv)
+			for (++argv; *argv; ++argv)
+			{
+				if (true && (indent += 1 + std::strlen(*argv)) > 75)
+				{
+					args << " ...";
+
+					break;
+				}
+
+				args << ' ' << *argv;
+			}
+
+		return args.str();
+	}
+};
+
+typedef kinfo_proc *InfoProc;
+
 int main(int argc, char *argv[])
 {
 	char error[_POSIX2_LINE_MAX];
-	kvm_t *kd = kvm_openfiles(NULL, _PATH_DEVNULL, NULL, O_RDONLY, error);
+	kvm_t *kd(kvm_openfiles(NULL, _PATH_DEVNULL, NULL, O_RDONLY, error));
 
 	if (!kd)
 		errx(1, "%s", error);
 
 	int count;
-	struct kinfo_proc *procs = kvm_getprocs(kd, KERN_PROC_PROC, 0, &count);
+	InfoProc procs(kvm_getprocs(kd, KERN_PROC_PROC, 0, &count));
 
 	if (!procs)
 		errx(1, "%s", kvm_geterr(kd));
 
-	for (struct kinfo_proc *proc = procs, *end = procs + count; proc != end; ++proc)
-		printf("%s %i %i\n", proc->ki_comm, proc->ki_pid, proc->ki_ppid);
+	PidMap pids;
+
+	_forall (InfoProc, proc, procs, procs + count)
+		pids[proc->ki_pid] = new Proc(kd, proc);
+
+	Sort sort(Name);
+
+	_foreach (PidMap, pid, pids)
+	{
+		Proc *proc(pid->second);
+		PidMap::iterator parent(pids.find(proc->parent()));
+
+		if (parent != pids.end())
+			parent->second->child(proc);
+	}
+
+	if (true)
+	{
+		pid_t hpid(0);
+		PidMap::iterator pid(pids.find(hpid ? hpid : getpid()));
+
+		if (pid != pids.end())
+			pid->second->highlight();
+	}
+
+	NameMap names;
+
+	_foreach (PidMap, pid, pids)
+	{
+		Proc *proc(pid->second);
+
+		if (proc->root())
+			switch (sort)
+			{
+			case Pid:
+				proc->printByPid();
+				break;
+			case Name:
+				names.insert(NameMap::value_type(proc->name(), proc));
+			}
+	}
+
+	switch (sort)
+	{
+	case Name:
+		_foreach (NameMap, name, names)
+			name->second->printByName();
+	default:
+		break;
+	}
 
 	return 0;
 }