/* pstree.c  -  display process tree. Written by Werner Almesberger */


#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <unistd.h>
#include <getopt.h>
#include <pwd.h>
#include <dirent.h>
#include <termios.h>
#include <termcap.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <linux/sched.h>


#define COMM_LEN     sizeof(dummy.comm)
#ifndef MAX_DEPTH
#define MAX_DEPTH    100
#endif
#define PROC_BASE    "/proc"


static struct task_struct dummy;

typedef struct _proc {
    char comm[COMM_LEN+1];
    pid_t pid;
    uid_t uid;
    int highlight;
    struct _child *children;
    struct _proc *parent;
    struct _proc *next;
} PROC;

typedef struct _child {
    PROC *child;
    struct _child *next;
}  CHILD;


static PROC *list = NULL;
static int width[MAX_DEPTH],more[MAX_DEPTH];
static int compact = 1,user_change = 0,pids = 0,trunc = 1;
static output_width = 132;
static int cur_x = 1;
static char last_char = 0;
static int dumped = 0; /* used by dump_by_user */


static void out_char(char c)
{
    if (cur_x < output_width || !trunc) putchar(c);
    if (cur_x == output_width && trunc)
	if (last_char) putchar('+');
	else {
	    last_char = c;
	    return;
	}
    cur_x++;
}


static void out_string(char *str)
{
    while (*str) out_char(*str++);
}


static int out_int(int x) /* non-negative integers only */
{
    int digits,div;

    digits = 0;
    for (div = 1; x/div; div *= 10) digits++;
    if (!digits) digits = 1;
    for (div /= 10; div; div /= 10) out_char('0'+(x/div) % 10);
    return digits;
}


static void out_newline(void)
{
    if (last_char && cur_x == output_width) putchar(last_char);
    last_char = 0;
    putchar('\n');
    cur_x = 1;
}


static PROC *find_proc(pid_t pid)
{
    PROC *walk;

    for (walk = list; walk; walk = walk->next)
	if (walk->pid == pid) break;
    return walk;
}


static PROC *new_proc(char *comm,pid_t pid,uid_t uid)
{
    PROC *new;

    if (!(new = malloc(sizeof(PROC)))) {
	perror("malloc");
	exit(1);
    }
    strcpy(new->comm,comm);
    new->pid = pid;
    new->uid = uid;
    new->highlight = 0;
    new->children = NULL;
    new->parent = NULL;
    new->next = list;
    return list = new;
}


static void add_child(PROC *parent,PROC *child)
{
    CHILD *new,**walk;
    int cmp;

    if (!(new = malloc(sizeof(CHILD)))) {
	perror("malloc");
	exit(1);
    }
    new->child = child;
    for (walk = &parent->children; *walk; walk = &(*walk)->next)
	if ((cmp = strcmp((*walk)->child->comm,child->comm)) > 0) break;
	else if (!cmp && (*walk)->child->uid > child->uid) break;
    new->next = *walk;
    *walk = new;
}


static void add_proc(char *comm,pid_t pid,pid_t ppid,uid_t uid)
{
    PROC *this,*parent;

    if (!(this = find_proc(pid))) this = new_proc(comm,pid,uid);
    else {
	strcpy(this->comm,comm);
	this->uid = uid;
    }
    if (!(parent = find_proc(ppid))) parent = new_proc("?",ppid,0);
    add_child(parent,this);
    this->parent = parent;
}


static int tree_equal(PROC *a,PROC *b)
{
    CHILD *walk_a,*walk_b;

    if (strcmp(a->comm,b->comm)) return 0;
    if (user_change && a->uid != b->uid) return 0;
    for (walk_a = a->children, walk_b = b->children; walk_a && walk_b;
      walk_a = walk_a->next, walk_b = walk_b->next)
	if (!tree_equal(walk_a->child,walk_b->child)) return 0;
    return !(walk_a || walk_b);
}


static void dump_tree(PROC *current,int level,int rep,int leaf,int last,
  uid_t prev_uid,int closing)
{
    CHILD *walk,*next,**scan;
    struct passwd *pw;
    int lvl,i,add,offset,info,count,comm_len;
    char *tmp,*here;
    char comm_tmp[5];

    if (!current) return;
    if (level >= MAX_DEPTH-1) {
	fprintf(stderr,"MAX_DEPTH not big enough.\n");
	exit(1);
    }
    if (!leaf)
	for (lvl = 0; lvl < level; lvl++) {
	    for (i = width[lvl]; i; i--) out_char(' ');
	    out_char(' ');
	    out_string(lvl == level-1 ? last ? "\\-" : "|-" : more[lvl+1] ?
	      "| " : "  ");
	}
    if (rep < 2) add = 0;
    else {
	add = out_int(rep);
	out_string("*[");
    }
    if (current->highlight && (tmp = tgetstr("md",NULL))) printf("%s",tmp);
    comm_len = 0;
    for (here = current->comm; *here; here++)
	if (*here == '\\') {
	    out_string("\\\\");
	    comm_len += 2;
	}
	else if (*here > ' ' && *here <= '~') {
		out_char(*here);
		comm_len++;
	    }
	    else {
		sprintf(comm_tmp,"\\%03o",(unsigned char) *here);
		out_string(comm_tmp);
		comm_len += 4;
	    }
    offset = cur_x;
    info = pids || (user_change && prev_uid != current->uid);
    if (info) out_char('(');
    if (pids) (void) out_int(current->pid);
    if (user_change && prev_uid != current->uid) {
	if (pids) out_char(',');
	if ((pw = getpwuid(current->uid))) out_string(pw->pw_name);
	else (void) out_int(current->uid);
    }
    if (info) out_char(')');
    if (current->highlight && (tmp = tgetstr("me",NULL))) printf("%s",tmp);
    if (!current->children) {
	while (closing--) out_char(']');
	out_newline();
    }
    else {
	more[level] = !last;
	width[level] = comm_len+cur_x-offset+add;
	out_string("---");
	if (cur_x >= output_width && trunc) out_char('+');
	else for (walk = current->children; walk; walk = next) {
		count = 0;
		next = walk->next;
		if (compact) {
		    scan = &walk->next;
		    while (*scan)
			if (!tree_equal(walk->child,(*scan)->child)) 
			    scan = &(*scan)->next;
			else {
			    if (next == *scan) next = (*scan)->next;
			    count++;
			    *scan = (*scan)->next;
			}
		}
		dump_tree(walk->child,level+1,count+1,walk == current->children,
		  !next,current->uid,closing+(count ? 1 : 0));
	    }
    }
}


static void dump_by_user(PROC *current,uid_t uid)
{
    CHILD *walk;

    if (current->uid == uid) {
	if (dumped) putchar('\n');
	dump_tree(current,0,1,1,1,uid,0);
	dumped = 1;
	return;
    }
    for (walk = current->children; walk; walk = walk->next)
	dump_by_user(walk->child,uid);
}


static void read_proc(void)
{
    DIR *dir;
    struct dirent *de;
    FILE *file;
    struct stat st;
    char path[PATH_MAX+1],comm[COMM_LEN+1];
    pid_t pid,ppid;
    int dummy;

    if (!(dir = opendir(PROC_BASE))) {
	perror(PROC_BASE);
	exit(1);
    }
    while (de = readdir(dir))
	if (pid = atoi(de->d_name)) {
	    sprintf(path,"%s/%d/stat",PROC_BASE,pid);
	    if (file = fopen(path,"r")) {
		if (fstat(fileno(file),&st) < 0) {
		    perror(path);
		    exit(1);
		}
		if (fscanf(file,"%d (%[^)]) %c %d",&dummy,comm,(char *) &dummy,
		  &ppid) == 4) add_proc(comm,pid,ppid,st.st_uid);
		(void) fclose(file);
	    }
	}
    (void) closedir(dir);
}


#if 0

/* Could use output of  ps achlx | awk '{ print $3,$4,$2,$13 }'  */

static void read_stdin(void)
{
    char comm[PATH_MAX+1];
    char *cmd;
    int pid,ppid,uid;

    while (scanf("%d %d %d %s\n",&pid,&ppid,&uid,comm) == 4) {
	if (cmd = strrchr(comm,'/')) cmd++;
	else cmd = comm;
	if (*cmd == '-') cmd++;
	add_proc(cmd,pid,ppid,uid);
    }
}

#endif


static void usage(char *name)
{
    fprintf(stderr,"usage: %s [ -c ] [ -h ] [ -l ] [ -p ] [ -u ] [ pid | user]"
      "\n\n",name);
    fprintf(stderr,"    -c     don't compact identical subtrees\n");
    fprintf(stderr,"    -h     highlight current process and its ancestors\n");
    fprintf(stderr,"    -l     don't truncate long lines\n");
    fprintf(stderr,"    -p     show PIDs; implies -c\n");
    fprintf(stderr,"    -u     show uid transitions\n");
    fprintf(stderr,"    pid    start at pid, default 1 (init)\n");
    fprintf(stderr,"    user   show only trees rooted at processes of that "
      "user\n\n");
    exit(1);
}


int main(int argc,char **argv)
{
    PROC *current;
    struct winsize winsz;
    struct passwd *pw;
    pid_t pid,highlight;
    int c;

    if (ioctl(1,TIOCGWINSZ,&winsz) >= 0)
	if (winsz.ws_col) output_width = winsz.ws_col;
    pid = 1;
    highlight = 0;
    pw = NULL;
    while ((c = getopt(argc,argv,"chplu")) != EOF)
	switch (c) {
	    case 'c':
		compact = 0;
		break;
	    case 'h':
		if (getenv("TERM") && tgetent(NULL,getenv("TERM")) == 1)
		    highlight = getpid();
		break;
	    case 'l':
		trunc = 0;
		break;
	    case 'p':
		pids = 1;
		compact = 0;
		break;
	    case 'u':
		user_change = 1;
		break;
	    default:
		usage(argv[0]);
	}
    if (optind == argc-1)
	if (isdigit(*argv[optind])) {
	    if (!(pid = atoi(argv[optind++]))) usage(argv[0]);
	}
	else if (!(pw = getpwnam(argv[optind++]))) {
		fprintf(stderr,"No such user name: %s\n",argv[optind-1]);
		return 1;
	    }
    if (optind != argc || (compact && pids)) usage(argv[0]);
    read_proc();
    for (current = find_proc(highlight); current; current = current->parent)
	current->highlight = 1;
    if (!pw)
	dump_tree(find_proc(pid),0,1,1,1,0,0);
    else {
	dump_by_user(find_proc(1),pw->pw_uid);
	if (!dumped) {
	    fprintf(stderr,"No processes found.\n");
	    return 1;
	}
    }
    return 0;
}
