/*
 * A very fast ps that goes through procfs
 * @(#)qps.c	1.21	08/02/97		Doug Hughes Auburn University
 *
 */
#include <stdio.h>
#include <sys/types.h>
#include <sys/fcntl.h>
#include <sys/procfs.h>
#include <string.h>
#include <stdlib.h>
#include <pwd.h>
#include <dirent.h>
#include <sys/mkdev.h>

#ifndef FALSE
#define FALSE 0
#endif
#ifndef TRUE
#define TRUE 1
#endif

static int reverse=0;

typedef struct qps {
	int pid;
	char ppid[40];
	char tty[3];
	char user[9];
	float cpu;
	float mem;
	char pname[17];
	char *arglist;
} *QPS;
	
void usage() 
{
	puts("qps [-acdmnptruPU]");
	puts("  -a  show initial (80 chars of) process args");
	puts("  -A  show ALL arguments");
	puts("  -c  sort by CPU usage");
	puts("  -d  debug (will print process ID before opening process)");
	puts("  -e  display environment variables");
	puts("  -m  sort by memory usage");
	puts("  -n  sort by process name");
	puts("  -p  sort py pid");
	puts("  -P  show parent process and PGID");
	puts("  -t  sort by tty");
	puts("  -r  reverse order sorting");
	puts("  -s  memory size info SIZE and RSS in pages and bytes (respectively)");
	puts("  -u  sort by user");
	puts("  -U  do not attempt to do uid->username mapping (Save CPU)");
	exit(1);
}

int sort_by_pid(const void *farg1, const void  *farg2)
{
	QPS arg1, arg2;

	arg1 = (QPS) farg1;
	arg2 = (QPS) farg2;
	if (reverse)
		return(arg2->pid > arg1->pid);
	else
		return(arg1->pid > arg2->pid);
}

int sort_by_cpu(const void *farg1, const void *farg2)
{
	QPS arg1, arg2;

	arg1 = (QPS) farg1;
	arg2 = (QPS) farg2;
	if (reverse)
		return(arg2->cpu < arg1->cpu);
	else
		return(arg1->cpu < arg2->cpu);
}

int sort_by_mem(const void *farg1, const void *farg2)
{
	QPS arg1, arg2;

	arg1 = (QPS) farg1;
	arg2 = (QPS) farg2;
	if (reverse)
		return(arg2->mem < arg1->mem);
	else
		return(arg1->mem < arg2->mem);
}

int sort_by_tty(const void *farg1, const void *farg2)
{
	QPS arg1, arg2;

	arg1 = (QPS) farg1;
	arg2 = (QPS) farg2;
	if (reverse)
		return(strcmp(arg2->tty, arg1->tty));
	else
		return(strcmp(arg1->tty, arg2->tty));
}

int sort_by_user(const void *farg1, const void *farg2)
{
	QPS arg1, arg2;

	arg1 = (QPS) farg1;
	arg2 = (QPS) farg2;
	if (reverse)
		return(strcmp(arg2->user, arg1->user));
	else
		return(strcmp(arg1->user, arg2->user));
}

int sort_by_pname(const void *farg1, const void *farg2)
{
	QPS arg1, arg2;

	arg1 = (QPS) farg1;
	arg2 = (QPS) farg2;
	if (reverse)
		return(strcmp(arg2->pname, arg1->pname));
	else
		return(strcmp(arg1->pname, arg2->pname));
}

/*
 * Subproc to display all args (as opposed to what's in arglist
 */
char *showallargs(int fd, struct prpsinfo *p)
{
	/* routine rewritten by Roger Faulkner - procfs co-designer and guru */
	char **argv;					/* hold arg pointers */
	char *env[400];					/* hold environment */
	int argc = p->pr_argc;			/* arg count */
	char buf[BUFSIZ+1];				/* temporary buffer */
	int len;						/* strlen(buf) */
	char *retbuf= NULL;				/* to return arglist */
	int argsize=64;					/* keep argsize for allocation records */
	int retbuflen;					/* strlen(retbuf) */
	int i = 0;
	char *sp;							/* a space character */


	if ((retbuf = malloc(argsize)) == NULL) {
		fprintf(stderr, "Out of memory in malloc!");
		exit(1);
	}
	retbuf[0] = (char) NULL;
	retbuflen = 0;
	buf[BUFSIZ] = (char) NULL;

	if (argc * sizeof(*argv) <= sizeof(env))
		argv = env;		/* use local buffer */
	else if ((argv = malloc(argc * sizeof(*argv))) == NULL) {
		fprintf(stderr, "Out of memory in malloc!");
		exit(1);
	}
	if ((argc = pread(fd, argv, argc * sizeof(*argv), p->pr_argv)) > 0)
		argc /= sizeof(*argv);

	for (i = 0; i < argc; i++) {
		if (argv[i] == NULL || pread(fd, buf, BUFSIZ, argv[i]) <= 0)
			continue;
		len = strlen(buf);

		/* If it's a bunch of spaces (zero'd out by process), skip/trunc it*/
		sp = strchr(buf, ' ');
		if (sp == buf) 
			continue;
		else if (sp != (char *) NULL) {
			sp = buf + len;
			while (*--sp == ' ') {
				*sp = (char) NULL;
				len--;
			}
		}

		while (retbuflen + len + 1 >= argsize) {
			argsize *= 2;
			if ((retbuf = realloc(retbuf, argsize)) == NULL) {
				fprintf(stderr, "Out of memory in realloc!");
				exit(1);
			}
		}
		if (retbuflen != 0)
			retbuf[retbuflen++] = ' ';
		strcpy(retbuf + retbuflen, buf);
		retbuflen += len;
	}

	if (argv != env)		/* if we allocated argv */
		free(argv);			/* free it */

	return(retbuf);
}
	
/*
 * subproc to display environment variables
 */

showenv(int fd, struct prpsinfo *p)
{
	char *env[400];					/* hold environment */
	char buf[BUFSIZ];				/* temporary buffer */
	int i = 0;

	pread(fd, env, 400, p->pr_envp);
	while (env[i] != NULL && i < 400) {
		pread(fd, buf, BUFSIZ, env[i++]);
		printf("%s ", buf);
	}
	puts("");
}

/*
 * Main program
 *
 * fast - regular qps, and qps -d (to find stuck processes waiting on I/O)
 *          qps -a (showargs), 
 *          qps -U (don't do NIS/NIS+ lookups),
 *          qps -s (show RSS, SIZE stuff)
 *        
 * medium - qps -e (has to munge through and grab out environment stuff.
 *			uses read on open descriptor and has to do pointers to pointers)
 *
 * slower - sorting of any kind. May get stuck if a process is hung waiting
 *           for I/O. If so, run qps -d. (It's still pretty darned fast)
 *          
 */

main(int argc, char *argv[])
{
	struct prpsinfo p;				/* process information structure */
	struct passwd *pw;				/* hold password lookup information */
	char rssinfo[80];				/* character resident set/memory info */
	DIR *dirf;						/* directory file pointer */
	char nothing[2]=" ";			/* Just empty printing stuff */
	struct dirent *dirp;			/* directory pointer */
	int fd;							/* file descriptor */
	char c;							/* for getopt */
	int cnt, i;						/* counting and looping */
	int mapuid = TRUE;				/* Map userid to username? */
	register int min, maj;			/* Major and minor device numbers */
	int (*func)(const void *, const void *) = NULL;		/* sorting function pointer */
	struct qps pstruct[1024];		/* Process structure */
	char parg[15];					/* process name buffer */
	char *ppstr = "";				/* parent and pgid */
	char *pphdr = " PPID  PGID ";	/* PPID and PGID headers */
	char *ppdelim = "";				/* delimiters for PPID and PGID */
	char *ppdelimd="----- ----- ";	
	char ppbuf[40] = "";			/* hold actual ppid and pgid */

	/* flags */
	int debug = 0;					/* debugging flag */
	int showargs = 0;				/* Show arguments */
	int rss = 0;					/* Resident set size */
	int environ = 0;				/* show environment? */
	int showppid = 0;				/* show parent and pgid */
	int Allargs = 0;				/* show ALL arguments */

	while ((c = getopt(argc, argv, "AadertuUpPcmns")) != -1) {
		switch(c) {
			case 'A':
				if (rss) {
					puts("rss and argument showing are mutually exclusive");
					exit(1);
				}
				Allargs = 1;
				break;
			case 'a':
				if (rss) {
					puts("rss and argument showing are mutually exclusive");
					exit(1);
				}
				showargs = 1;
				break;
			case 'c':
				func = &sort_by_cpu;
				break;
			case 'd':
				debug = 1;
				break;
			case 'e':
				environ = 1;
				break;
			case 'm':
				func = &sort_by_mem;
				break;
			case 'n':
				func = &sort_by_pname;
				break;
			case 'p':
				func = &sort_by_pid;
				break;
			case 'P':
				showppid = 1;
				ppdelim = ppdelimd;
				ppstr = pphdr;
				break;
			case 'r':
				reverse = 1;
				break;
			case 's':
				if (showargs) {
					puts("rss and argument showing are mutually exclusive");
					exit(1);
				}
				rss = 1;
				break;
			case 't':
				func = &sort_by_tty;
				break;
			case 'u':
				func = &sort_by_user;
				break;
			case 'U':
				mapuid = FALSE;
				break;
			default:
				usage();
				break;
		}
	}

	/* Set line buffering for output */
	setvbuf(stdout, NULL, _IOLBF, 1024);

	if ((dirf = opendir("/proc")) == NULL) {
		perror("couldn't open proc");
		exit(1);
	}
	(void) readdir(dirf); /* skip over . and .. */
	(void) readdir(dirf);
	
	/* Display headers */
	printf("%-7s %s%-3.3s %-8.8s %-4.4s %-4.4s %-17.17s", "PID", ppstr, "TTY", "User", "%CPU", "%Mem", "Process Name");
	if (rss) /* print process size headers */
		printf(" %5s %5s %6s %6s\n", "SIZEP", "RSSP", "SIZE_B", "RSS_B");
	else
		printf("\n");
		
	printf("%-7s %s%-3.3s %-8.8s %-4.4s %-4.4s %-17.17s", "------", ppdelim, "---", "-------", "----", "----", "-----------------");
	if (rss) /* print process size headers */
		printf(" %5s %5s %6s %6s\n", "-----", "-----", "------", "------");
	else
		printf("\n");

	cnt = 0;
	/* open /procfs and scan through files one at a time - each a process */
	while ((dirp = readdir(dirf)) != NULL)  {
		sprintf(parg, "/proc/%s", dirp->d_name);
	
		if (debug)
			printf("process %s\n", dirp->d_name);

		if ((fd = open(parg, O_RDONLY)) < 0) {
			continue;
		}

		/* Grab process information/status */
		if (ioctl(fd, PIOCPSINFO, (void *) &p) < 0) {
			close(fd);
			continue;
		}
	
		/* map major and minor device numbers */
		min = minor(p.pr_lttydev);
		maj = major(p.pr_lttydev);

		/* Show process arguments with name */
		if (Allargs)  {
			pstruct[cnt].arglist = showallargs(fd, &p);
		} else if (showargs) {
			pstruct[cnt].arglist = strdup(p.pr_psargs);
		} else if (rss) {
			sprintf(rssinfo, " %5d %5d %6u %6u", p.pr_size, p.pr_rssize,
				p.pr_bysize/1024, p.pr_byrssize/1024);
			pstruct[cnt].arglist = strdup(rssinfo);
		} else {
			pstruct[cnt].arglist = nothing;
		}
	
		/* Do we want to see parent and pgrp? */
		if (showppid)
			sprintf(ppbuf, "%5d %5d ", p.pr_ppid, p.pr_pgrp);

		/* Convert mem and CPU usage to percentage of machine capacity */
		pstruct[cnt].pid = p.pr_pid;
		strcpy(pstruct[cnt].ppid, ppbuf);
		pstruct[cnt].mem = p.pr_pctmem * 100.0 / (float) 0x8000;
		pstruct[cnt].cpu = p.pr_pctcpu * 100.0 / (float) 0x8000;
		/* Get uid/uname */
		if (!mapuid || (pw = getpwuid(p.pr_uid)) == NULL) 
			sprintf(pstruct[cnt].user, "%-8d", p.pr_uid);
		else
			sprintf(pstruct[cnt].user, "%-8.8s", pw->pw_name);
		sprintf(pstruct[cnt].pname, "%-16.16s", p.pr_fname);
		/* check controlling tty if available */
		if (p.pr_ttydev == PRNODEV)
			sprintf(pstruct[cnt].tty, "--");
		else if (maj == 0)
			sprintf(pstruct[cnt].tty, "co");
		else
			sprintf(pstruct[cnt].tty, "%-.1d", min);
		if (func == NULL) {	/* no sorting - just do it */
			printf("%-7d %s%-3.2s %-8.8s %4.1f %4.1f %-16.16s %s\n", pstruct[cnt].pid, pstruct[cnt].ppid, pstruct[cnt].tty, pstruct[cnt].user, pstruct[cnt].cpu, pstruct[cnt].mem, pstruct[cnt].pname, pstruct[cnt].arglist);
			if (environ)
				showenv(fd, &p);
		}
		cnt++;
		close(fd);
	}
	closedir(dirf);

	/* sorting applied */
	if (func != NULL) {
		qsort((void*) pstruct, cnt, sizeof(struct qps), func);
	
		for (i=0; i < cnt; i++) 
			printf("%-7d %s%-3.2s %-8.8s %4.1f %4.1f %-16.16s %s\n", pstruct[i].pid, pstruct[i].ppid, pstruct[i].tty, pstruct[i].user, pstruct[i].cpu, pstruct[i].mem, pstruct[i].pname, pstruct[i].arglist );
	}

	exit(0);
}
