#!/usr/bin/perl
use strict;
use Getopt::Long;

my ($threads, $kthreads, $callmode, @uids, $vserver, $xid, $showxid, $help);
my (%_uids, $_sockets);

my @SYSCALLS = qw(
	read write open close stat fstat lstat poll lseek mmap mprotect munmap
	brk rt_sigaction rt_sigprocmask rt_sigreturn ioctl pread64 pwrite64
	readv writev access pipe select sched_yield mremap msync mincore
	madvise shmget shmat shmctl dup dup2 pause nanosleep getitimer alarm
	setitimer getpid sendfile socket connect accept sendto recvfrom sendmsg
	recvmsg shutdown bind listen getsockname getpeername socketpair
	setsockopt getsockopt clone fork vfork execve exit wait4 kill uname
	semget semop semctl shmdt msgget msgsnd msgrcv msgctl fcntl flock fsync
	fdatasync truncate ftruncate getdents getcwd chdir fchdir rename mkdir
	rmdir creat link unlink symlink readlink chmod fchmod chown fchown
	lchown umask gettimeofday getrlimit getrusage sysinfo times ptrace
	getuid syslog getgid setuid setgid geteuid getegid setpgid getppid
	getpgrp setsid setreuid setregid getgroups setgroups setresuid
	getresuid setresgid getresgid getpgid setfsuid setfsgid getsid capget
	capset rt_sigpending rt_sigtimedwait rt_sigqueueinfo rt_sigsuspend
	sigaltstack utime mknod uselib personality ustat statfs fstatfs sysfs
	getpriority setpriority sched_setparam sched_getparam
	sched_setscheduler sched_getscheduler sched_get_priority_max
	sched_get_priority_min sched_rr_get_interval mlock munlock mlockall
	munlockall vhangup modify_ldt pivot_root _sysctl prctl arch_prctl
	adjtimex setrlimit chroot sync acct settimeofday mount umount2 swapon
	swapoff reboot sethostname setdomainname iopl ioperm create_module
	init_module delete_module get_kernel_syms query_module quotactl
	nfsservctl getpmsg putpmsg afs_syscall tuxcall security gettid
	readahead setxattr lsetxattr fsetxattr getxattr lgetxattr fgetxattr
	listxattr llistxattr flistxattr removexattr lremovexattr fremovexattr
	tkill time futex sched_setaffinity sched_getaffinity set_thread_area
	io_setup io_destroy io_getevents io_submit io_cancel get_thread_area
	lookup_dcookie epoll_create epoll_ctl_old epoll_wait_old
	remap_file_pages getdents64 set_tid_address restart_syscall semtimedop
	fadvise64 timer_create timer_settime timer_gettime timer_getoverrun
	timer_delete clock_settime clock_gettime clock_getres clock_nanosleep
	exit_group epoll_wait epoll_ctl tgkill utimes vserver mbind
	set_mempolicy get_mempolicy mq_open mq_unlink mq_timedsend
	mq_timedreceive mq_notify mq_getsetattr kexec_load waitid add_key
	request_key keyctl ioprio_set ioprio_get inotify_init inotify_add_watch
	inotify_rm_watch migrate_pages openat mkdirat mknodat fchownat
	futimesat newfstatat unlinkat renameat linkat symlinkat readlinkat
	fchmodat faccessat pselect6 ppoll unshare set_robust_list
	get_robust_list splice tee sync_file_range vmsplice move_pages
);

my %fdops = (
	read => 0,
	write => 0,
	flock => 0,
	accept => 0,
	ioctl => 0,
);

my %intargs = (
	wait4 => 0,
);

##############################################################################

sub slurp {
	my ($file) = @_;
	my $text = "";
	open FH, "<", $file or return;
	$text .= $_ while <FH>;
	close FH;
	chomp $text;
	return $text;
}

sub Hex2IP {
	my $ip = hex(shift);
	return sprintf("%d.%d.%d.%d",
		$ip & 255,
		($ip >> 8) & 255,
		($ip >> 16) & 255,
		($ip >> 24) & 255
	);
}

sub U2S64 {
	my ($arg) = @_;
	my $shift = 1<<63;
	if($arg >= $shift) {
		$arg = ($arg & ~$shift) - $shift;
	}
	return $arg;
}

sub LoadSockets {
	$_sockets = {};
	open TCP, "</proc/net/tcp";
	while(my $line = <TCP>) {
		chomp $line;
		my ($sl, $lhost, $lport, $rhost, $rport, $state,
			$txq, $rxq, $tr, $tm, $rtt, $uid, $timeout, $inode) =
		$line =~ m{
			^\s*(\d+):\s
			(\w{8}):(\w{4})\s   # local
			(\w{8}):(\w{4})\s   # remote
			(\w\w)\s            # state
			(\w{8}):(\w{8})\s   # tx/rx queue
			(\w\w):(\w{8})\s+   # tr, tm->when
			(\d+)\s+            # retransmit
			(\d+)\s+            # uid
			(\d+)\s+            # timeout
			(\d+)\s+            # inode
			}x or next;
		$_sockets->{$inode} = {
			local => Hex2IP($lhost) . ":" . hex($lport),
			remote => Hex2IP($rhost) . ":" . hex($rport),
		};
	}
	close TCP;
}

sub getcall {
	my ($pid, $tid) = @_;
	$tid ||= $pid;

	if($callmode eq 'syscall') {
		my $sc = slurp("/proc/$pid/task/$tid/syscall") or return;
		if(my ($call, @args) = $sc =~ /^(\d+) 0x(\w+) 0x(\w+) 0x(\w+) 0x(\w+) 0x(\w+) 0x(\w+) 0x(\w+) 0x(\w+)$/) {
			@args = map { hex($_) } @args;
			my $callname = $SYSCALLS[$call];
			if(defined $fdops{$callname}) {
				my $fd = $args[$fdops{$callname}];
				my $file = readlink("/proc/$pid/fd/$fd");
				if($file) {
					if(my ($inode) = $file =~ /^socket:\[(\d+)\]$/) {
						LoadSockets() unless defined $_sockets;
						if(my $sock = $_sockets->{$1}) {
							$file = ($sc eq 'accept') ? $sock->{local} : $sock->{remote};
						}
					}
					return "$callname($file)";
				} else {
					return "$callname(fd $fd)";
				}
			} elsif(defined $intargs{$callname}) {
				my $int = U2S64($args[$intargs{$callname}]);
				return "$callname($int)";
			} else {
				return $callname;
			}
		} elsif(my ($sp, $pc) = $sc =~ /^-1 0x(\w+) 0x(\w+)$/) {
			return "0x$pc";
		} else {
			return "?";
		}
	} elsif($callmode eq 'wchan') {
		my $sc = slurp("/proc/$pid/task/$tid/wchan");
		$sc = "?" if $sc eq '_stext'; # this isn't useful
		return $sc;
	} else {
		return "";
	}
}

##############################################################################

my @ORIG_ARGV = @ARGV;
Getopt::Long::GetOptions(
		"user=s@"    => \@uids,
		"kthreads!"  => \$kthreads,
		"syscall!"   => sub { $callmode = $_[1] ? "syscall" : "" },
		"wchan!"     => sub { $callmode = $_[1] ? "wchan" : "" },
		"threads!"   => \$threads,
		"vserver=s"  => \$vserver,
		"xid=i"      => \$xid,      # internal
		"help"       => \$help,
		);

my $COLS = `tput cols` || 80;

# Do we actually have /proc/pid/syscall support?
my $have_syscall = -f "/proc/self/syscall";

# --syscall requires support, and overrides wchan
$callmode = ($have_syscall ? "syscall" : "wchan") unless defined $callmode;
$callmode = "wchan" if $callmode eq "syscall" && !$have_syscall;

# Put together the format...
my $pidwidth = 5;
my $userwidth = 10;
my $tcwidth = $callmode ? 30 : 0;
my $memwidth = 7; # hard
my $xidwidth = defined $xid ? 6 : 0;
my $fmt = "%${pidwidth}s %-${userwidth}s %-${tcwidth}s %${memwidth}s %s\n";

if(defined $help) {
	print <<EOF;
Usage: wps [OPTION]...
List processes with syscall/wchan information.

  -u, --user=USER            only display processes owned by this user.
                             may be repeated for multiple users.
  -k, --[no-]kthreads        display kernel threads.
  -t, --[no-]threads         display individual threads for processes.
  -s, --[no-]syscall         display syscalls for threads/processes.
  -w, --[no-]wchan           display wchan for threads/processes.
  -v, --vserver=VSERVER      display vserver processes.
                             VSERVER may be a PS name, XID, or "all".

EOF
	require POSIX;
	my ($sname, $nname, $rel, $ver, $mach) = POSIX::uname();
	if($have_syscall) {
		print "This kernel ($sname $rel) supports syscall display.\n";
	} else {
		print "This kernel ($sname $rel) doesn't support syscall display. Lame!\n";
	}
	exit 0;
}

if(defined $vserver && !defined $xid) {
	if($vserver eq 'all') {
		$xid = 1;
	} elsif(-e "/etc/vservers/$vserver/context") {
		$xid = slurp("/etc/vservers/$vserver/context");
	} elsif($vserver =~ /^\d+$/) {
		$xid = $vserver;
	} else {
		die "no such vserver '$vserver'\n";
	}
	my @args = (qw( /usr/sbin/vcontext --migrate --xid 1 -- /usr/bin/perl ), $0, @ORIG_ARGV, "--xid", $xid);
	exec @args;
	die "vcontext: $!";
}

@uids = map {
	scalar(getpwnam($_))
} @uids;

opendir PROC, "/proc" or die "opendir: $!";
printf $fmt, "PID", "USER", $callmode ? "CALL" : "", "MEM", "COMMAND";
while(my $pid = readdir PROC) {
	next if $pid =~ /\D/;

	my $pxid;
	if(defined $xid) {
		($pxid) = slurp("/proc/$pid/vinfo") =~ /^XID:\s+(\d+)$/m;
		next if $pxid != $xid && $xid != 1;
	}

	my $uid = (stat "/proc/$pid")[4];
	next if @uids && !grep { $_ eq $uid } @uids;
	my ($user) = $_uids{$uid} || ($_uids{$uid} = (getpwuid($uid) || $uid));

	my $cmd = slurp("/proc/$pid/cmdline");
	if($cmd eq '') {
		next unless $kthreads;
		($cmd) = slurp("/proc/$pid/stat") =~ /\((.+)\)/;
		$cmd ||= "unknown";
		$cmd = "[$cmd]";
	} else {
		$cmd =~ s/\0/ /g;
		$cmd =~ s/[^ -~]/?/g;
	}

	my %tcinfo;
	if($threads) {
		opendir TASK, "/proc/$pid/task" or next;
		while(my $tid = readdir TASK) {
			next if $tid =~ /\D/;
			$tcinfo{$tid} = getcall($pid, $tid);
		}
		closedir TASK;
	} else {
		$tcinfo{$pid} = getcall($pid, $pid);
	}

	my $meminfo = slurp("/proc/$pid/status");
	my ($mem) = $meminfo =~ /^VmRSS:\s+(\d+) kB$/m;
	if($mem < 1024 * 1.5) {
		$mem = sprintf "%d kB", $mem;
	} elsif($mem < 1048576 * 1.5) {
		$mem = sprintf "%d MB", $mem / 1024;
	} else {
		$mem = sprintf "%.1f GB", $mem / 1048576;
	}

	$cmd = substr $cmd, 0, ($COLS - $tcwidth - $userwidth - $memwidth - $pidwidth - 4);

	$user = (length $user > $userwidth) ? $uid : $user;
	$user = ($pxid <= 1 ? "HOST" : "ps$pxid") if $xid == 1;

	printf $fmt, $pid, $user, substr($tcinfo{$pid}, 0, $tcwidth), $mem, $cmd;

	next unless $threads;
	for my $tid (keys %tcinfo) {
		printf $fmt, "", ' \-', substr($tcinfo{$tid}, 0, $tcwidth), "", "thread $tid" unless $tid == $pid;
	}
}

