/* $Id: ptree.c,v 2.22 2012/09/24 17:45:38 ksb Exp $
 * This emulates the Solairs "ptree" command, which outputs	(petef&ksb)
 * a graph of the processes (presently running on this host)
 * showing which proceses are parents of which.
 *
 * $Compile: ${cc-cc} %f -lkvm -o %F
 */
#include <sys/param.h>
#if !defined(HPUX11)
#include <sys/user.h>
#endif
#include <fcntl.h>
#include <limits.h>
#include <paths.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>
#include <unistd.h>
#include <ctype.h>
#include <sys/types.h>
#include <pwd.h>

#include "machine.h"
#include "sparse.h"

const char *progname = "ptree";
const char rcsid[] =
	"$Id: ptree.c,v 2.22 2012/09/24 17:45:38 ksb Exp $";

/* Output a helpful usage map						(ksb)
 */
static void
Usage(FILE *fp, int fVerbose)
{
	fprintf(fp, "%s: usage [-an] [pids|logins]\n", progname);
	fprintf(fp, "%s: usage -h\n", progname);
	fprintf(fp, "%s: usage -V\n", progname);
	if (!fVerbose)
		return;
	fprintf(fp, "a      include children of process 0\n");
	fprintf(fp, "h      display only this help text\n");
	fprintf(fp, "n      do not trace up towards init from target processes\n");
	fprintf(fp, "V      display only some version information\n");
	fprintf(fp, "pids   target child process of each pid\n");
	fprintf(fp, "logins target each login\'s processes\n");
}

/* pid order							(ksb)
 */
int
PidOrder(const void *pvLeft, const void *pvRight)
{
	return ((struct kinfo_proc **)pvLeft)[0]->PTREE_FIND_PID - ((struct kinfo_proc **)pvRight)[0]->PTREE_FIND_PID;
}

/* Establish the KVM interface					(petef,ksb)
 * open kvm interface so we can read the process table,
 * then read info about all of the process table we can see.
 */
static kvm_t *pkvm;		/* our view into the kvm system		*/
struct kinfo_proc *pKIAll;	/* a vector of all processes		*/
struct kinfo_proc **ppKIParent;	/* ... and their parent			*/
struct kinfo_proc ***pppKIKids;	/* ... and their kids			*/
int iEntries;			/* number of entries in pKIAll		*/

static struct kinfo_proc **
SetKVM()
{
	auto char acError[_POSIX2_LINE_MAX];
	const char *nlistf, *memf;
	struct kinfo_proc **ppKIMem, **ppKIRet;
	unsigned uScan;

#if USE_KVM_FROM_GETPROCS
	memf = nlistf = (char *)0;
#else
	memf = nlistf = _PATH_DEVNULL;
#endif
	if ((kvm_t *)0 == (pkvm = kvm_openfiles(nlistf, memf, (char *)0, O_RDONLY, acError))) {
		fprintf(stderr, "%s: kvm_openfiles: %s\n", progname, acError);
		exit(EX_OSERR);
	}
	if ((struct kinfo_proc *)0 == (pKIAll = kvm_getprocs(pkvm, KERN_PROC_ALL, 0, & iEntries))) {
		fprintf(stderr, "%s: kvm_getprocs: %s\n", progname, kvm_geterr(pkvm));
		exit(EX_OSERR);
	}

	/* make the data structure to traverse
	 */
	ppKIMem = (struct kinfo_proc **)calloc((2*iEntries|7)+1, sizeof(struct kinfo_proc *));
	ppKIParent = (struct kinfo_proc **)calloc(iEntries+1, sizeof(struct kinfo_proc **));
	pppKIKids = (struct kinfo_proc ***)calloc(iEntries+1, sizeof(struct kinfo_proc **));
	for (uScan = 0; uScan < iEntries; ++uScan) {
		register unsigned uKid;

		if (0 == pKIAll[uScan].PTREE_FIND_PID) {
			pppKIKids[uScan] = (struct kinfo_proc **)0;
			continue;
		}
		pppKIKids[uScan] = ppKIMem;
		for (uKid = 0; uKid < iEntries; ++uKid) {
			if (pKIAll[uScan].PTREE_FIND_PID != pKIAll[uKid].PTREE_FIND_PPID)
				continue;
			ppKIParent[uKid] = & pKIAll[uScan];
			*ppKIMem++ = & pKIAll[uKid];
		}
		qsort(pppKIKids[uScan], ppKIMem-pppKIKids[uScan], sizeof(struct kinfo_proc *), PidOrder);
		*ppKIMem++ = (struct kinfo_proc *)0;
	}

	/* Build a list of init-like process, but not kernel threads
	 */
	ppKIRet = ppKIMem;
	for (uScan = 0; uScan < iEntries; ++uScan) {
		if (0 == pKIAll[uScan].PTREE_FIND_PID || (struct kinfo_proc *)0 != ppKIParent[uScan]) {
			continue;
		}
		*ppKIMem++ = & pKIAll[uScan];
	}
	qsort(ppKIRet, ppKIMem-ppKIRet, sizeof(struct kinfo_proc *), PidOrder);
	*ppKIMem++ = (struct kinfo_proc *)0;
	return ppKIRet;
}

static void
	**ppvPids;	/* pids asked for			*/
static char
	acTrail[] = "h&g",	/* just print a trail		*/
	acMark[] = "special";	/* output all kids as well	*/

/* Output the proc tree from here down				(petef,ksb)
 */
static void
PrintProc(struct kinfo_proc *pKIProc, int iIndent, int fMust, int fCan)
{
	register int i, fLevel;
	register char **ppcCmdArgs;
	auto struct kinfo_proc **ppKIKids;

	fLevel = 0;
	if (! fMust && (void **)0 != ppvPids && sp_exists(ppvPids, pKIProc->PTREE_FIND_PID)) {
		register void **ppvThis;
		ppvThis = sp_index(ppvPids, pKIProc->PTREE_FIND_PID);
		if (acTrail == *ppvThis) {
			fLevel = 1;
		} else {
			fMust = 1;
		}
	}
	if (fLevel || fMust || (fCan && 0 != pKIProc->PTREE_FIND_PPID)) {
		printf("%*.*s%-5d", iIndent, iIndent, " ", pKIProc->PTREE_FIND_PID);
		ppcCmdArgs = kvm_getargv(pkvm, pKIProc, MAXCOMLEN);
		if ((char **)0 == ppcCmdArgs)
			printf(" (%s)", pKIProc->PTREE_FIND_COMM);
		else
			for (i = 0; (char *)0 != ppcCmdArgs[i]; ++i)
				printf(" %s", ppcCmdArgs[i]);
		printf("\n");
		iIndent += 2;
	}

	for (ppKIKids = pppKIKids[pKIProc-pKIAll]; (struct kinfo_proc *)0 != *ppKIKids; ++ppKIKids) {
		PrintProc(*ppKIKids, iIndent, fMust, fCan);
	}
}

static int fAll, fTraceBack, fQuiet;	/* -a, -n, -q */

/* Mark the pid as interesting up the tree to init			(ksb)
 */
static void
MarkTree(struct kinfo_proc *pKI)
{
	register char *pcMark;
	register void **ppvFind;

	for (pcMark = acMark; (struct kinfo_proc *)0 != pKI; pKI = ppKIParent[pKI-pKIAll]) {
		if ((void **)0 != (ppvFind = sp_index(ppvPids, pKI->PTREE_FIND_PID))) {
			*ppvFind = pcMark;
		}
		if (!fTraceBack || (!fAll && 1 == pKI->PTREE_FIND_PPID)) {
			break;
		}
		pcMark = acTrail;
	}
}

/* Process a few options, open the kvm interface, find the	(petef, ksb)
 * processes we were asked to report on, output the graph.
 */
int
main(int argc, char **argv)
{
	struct kinfo_proc **ppKIInits;
	int i, fGivenPid, fGivenUid;
	auto void **ppvUids;			/* any uids requested	*/

	if ((char *)0 != argv[0]) {
		if ((char *)0 != (progname = strrchr(argv[0], '/')))
			++progname;
		else
			progname = argv[0];
	}

	/* Find all the pids that look like init's children.		(ksb)
	 * We'll start there if we're looking to see the entire process
	 * table (no pid arg), or we filter on the sparse spaces.
	 */
	ppKIInits = SetKVM();
	if (0 == iEntries) {
		fprintf(stderr, "%s: can't find any processes\n", progname);
		exit(EX_TEMPFAIL);
	}

	/* parse request
	 */
	ppvPids = sp_init(~0);
	ppvUids = sp_init(~0);
	fQuiet = 0;
	fTraceBack = 1;
	fAll = 0;
	while (EOF != (i = getopt(argc, argv, "ahnVPq"))) {
		switch (i) {
		case 'a':
			fAll = 1;
			break;
		case 'h':
			Usage(stdout, 1);
			exit(EX_OK);
		case 'n':
			fTraceBack = 0;
			break;
		case 'q':
			fQuiet = 1;
			break;
		case 'V':
			printf("%s: %s\n", progname, rcsid);
			exit(EX_OK);
		case '?':
		default:
			Usage(stderr, 0);
			exit(EX_USAGE);
		}
	}
	argc -= optind;
	argv += optind;
	setpwent();
	for (fGivenUid = fGivenPid = 0; 0 < argc--; ++argv) {
		register struct passwd *pPW;
		register void **ppvLoved;
		register pid_t iPid;
		register struct kinfo_proc *pKI;

		if (isdigit((int)argv[0][0]) || '-' == argv[0][0]) {
			pKI = (struct kinfo_proc *)0;
			iPid = ('-' == argv[0][0] && '\000' == argv[1]) ? getpid() : atol(argv[0]);
			for (i = 0; i < iEntries; ++i) {
				if (iPid == pKIAll[i].PTREE_FIND_PID) {
					pKI = & pKIAll[i];
					break;
				}
			}

			if ((struct kinfo_proc *)0 == pKI) {
				fprintf(stderr, "%s: %s: no such process\n", progname, argv[0]);
				fQuiet = 1;
				*ppKIInits = (struct kinfo_proc *)0;
				continue;
			}
			++fGivenPid;
			MarkTree(pKI);
			continue;
		}
		if ((struct passwd *)0 == (pPW = getpwnam(argv[0]))) {
			fprintf(stderr, "%s: %s: argument not a valid login\n", progname, argv[0]);
			exit(EX_DATAERR);
		}
		if ((void **)0 != (ppvLoved = sp_index(ppvUids, pPW->pw_uid))) {
			*ppvLoved = acMark;
			++fGivenUid;
		}
	}
	endpwent();
	if (fGivenUid) {
		register int fFound = 0;

		for (i = 0; i < iEntries; ++i) {
			if (!sp_exists(ppvUids, pKIAll[i].PTREE_FIND_RUID)) {
				continue;
			}
			fGivenPid++;
			MarkTree(& pKIAll[i]);
			fFound = 1;
		}
		if (!fFound) {
			fQuiet = 1;
		}
	}
	if (!fGivenPid) {
		ppvPids = (void **)0;
	}

	/* When we have no idea where to start, force kids of init as a,
	 * start rather than init-like processes.
	 */
	if (fAll && (fGivenPid || fGivenUid)) {
		fAll = 0;
	}
	if (!fAll && !fGivenPid && !fGivenUid) {
		if (fQuiet) {
			exit(EX_OK);
		}
		for (i = 0; i < iEntries; ++i) {
			if (1 == pKIAll[i].PTREE_FIND_PID) {
				ppKIInits = pppKIKids[i];
				break;
			}
		}
		fAll = 1;
	}
	for (/*SetKVM*/; (struct kinfo_proc *)0 != *ppKIInits; ++ppKIInits) {
		PrintProc(*ppKIInits, 0, fAll, !(fGivenUid || fGivenPid));
	}
	exit(EX_OK);
}