diff --git a/pupy/modules/ps.py b/pupy/modules/ps.py index 60be0114..a60b9355 100644 --- a/pupy/modules/ps.py +++ b/pupy/modules/ps.py @@ -1,44 +1,192 @@ # -*- coding: utf-8 -*- from pupylib.PupyModule import * from pupylib.utils.rpyc_utils import obtain +from pupylib.utils.term import terminal_size, colorize from modules.lib.utils.shell_exec import shell_exec +import logging __class_name__="PsModule" +ADMINS = ('NT AUTHORITY\SYSTEM', 'root') + +def gen_columns(record, colinfo): + columns = {} + + columns['name'] = record.get('name') or '?' + columns['cmdline'] = ' '.join([ + x for x in record['cmdline'][1:] if x.strip() + ]) if 'cmdline' in record else '' + columns['exe'] = record.get('exe') or '{{{}}}'.format(columns['name']) + columns['username'] = record.get('username') or '' + cpu = record.get('cpu_percent') + columns['cpu_percent'] = '{:3}%'.format(int(cpu)) if cpu is not None else ' '*4 + mem = record.get('memory_percent') + columns['memory_percent'] = '{:3}%'.format(int(mem)) if mem is not None else ' '*4 + + if colinfo: + columns['username'] = '{{:{}}}'.format(colinfo['username']).format(columns['username']) + columns['pid'] = '{{:{}}}'.format(colinfo['pid']).format(record['pid']) + else: + columns['pid'] = '{}'.format(parent) + + return columns + +def gen_output_line(columns, info, record, width): + cpu = record.get('cpu_percent') or 0 + mem = record.get('memory_percent') or 0 + + if record.get('self'): + color = "green" + elif cpu > 70 or mem > 50: + color = "red" + elif record.get('username') in ADMINS: + if record.get('connections'): + color = "magenta" + else: + color = "yellow" + elif record.get('connections'): + color = "cyan" + elif not record.get('same_user'): + color = "grey" + else: + color = None + + template = ' '.join('{{{}}}'.format(x) for x in info) + output = template.format(**columns) + if width: + output = output[:width] + + if color: + output = colorize(output, color) + + return output + +def print_pstree(fout, parent, tree, data, + prefix='', indent='', width=80, colinfo={}, + info=['exe', 'cmdline'], hide=[], + first=False): + if parent in data: + data[parent]['pid'] = parent + columns = gen_columns(data[parent], colinfo) + + if ( columns['name'] in hide ) or ( columns['exe'] in hide ) or ( parent in hide ): + return + + columns['prefix'] = prefix + + before_tree = [ x for x in info if x in ('cpu_percent', 'memory_percent', 'username') ] + after_tree = [ x for x in info if x in ('exe', 'name', 'cmdline') ] + + outcols = [ 'pid' ] + before_tree + [ 'prefix' ] + after_tree + + fout.write(gen_output_line(columns, outcols, data[parent], width)+'\n') + + if parent not in tree: + return + + children = tree[parent][:-1] + + for child in children: + print_pstree( + fout, child, tree, data, + prefix=indent+('┌' if first else '├'), indent=indent + '│ ', width=width, + colinfo=colinfo, info=info, hide=hide + ) + + child = tree[parent][-1] + print_pstree( + fout, child, tree, data, + prefix=indent+'└', indent=indent + ' ', + width=width, colinfo=colinfo, + info=info, hide=hide + ) + +def print_ps(fout, data, width=80, colinfo={}, + info=['exe', 'cmdline'], hide=[]): + + outcols = [ 'pid' ] + [ + x for x in info if x in ('cpu_percent', 'memory_percent', 'username', 'exe', 'name', 'cmdline') + ] + + for process in sorted(data): + data[process]['pid'] = process + columns = gen_columns(data[process], colinfo) + + if ( columns['name'] in hide ) or ( columns['exe'] in hide ) or ( process in hide ): + continue + + fout.write(gen_output_line(columns, outcols, data[process], width)+'\n') + + @config(cat="admin") class PsModule(PupyModule): """ list processes """ - dependencies = { - 'windows': ['pupwinutils.processes'] - } + dependencies = [ 'pupyps' ] def init_argparse(self): self.arg_parser = PupyArgumentParser(prog="ps", description=self.__doc__) - self.arg_parser.add_argument('--all', '-a', action='store_true', help='more info') + self.arg_parser.add_argument('--tree', '-t', action='store_true', help='draw tree') + self.arg_parser.add_argument('-i', '--info', action='store_true', help='print more info') + self.arg_parser.add_argument('-x', '--hide', nargs='+', default=[], help='hide processes by pid/name/exe') + self.arg_parser.add_argument('-a', '--all', action='store_true', help='show kthread') + self.arg_parser.add_argument('-w', '--wide', action='store_true', help='show all arguments') + self.arg_parser.add_argument('-s', '--show', nargs='+', type=int, default=[], + help='show process (tree) by pid') def run(self, args): - if self.client.is_windows(): - outputlist=self.client.conn.modules["pupwinutils.processes"].enum_processes() - outputlist=obtain(outputlist) #pickle the list of proxy objects with obtain is really faster - columns=['username', 'pid', 'arch', 'exe'] - if args.all: - columns=['username', 'pid', 'arch', 'name', 'exe', 'cmdline', 'status'] - for dic in outputlist: - for c in columns: - if c in dic and dic[c] is None: - dic[c]="" - dic["cmdline"]=' '.join(dic['cmdline'][1:]) + width, _ = terminal_size() + root, tree, data = self.client.conn.modules.pupyps.pstree() + data = { int(k):v for k,v in obtain(data).iteritems() } + tree = { int(k):v for k,v in obtain(tree).iteritems() } + + colinfo = {'pid': 0} + for pid in data: + l = len(str(pid)) + if colinfo['pid'] < l: + colinfo['pid'] = l + for column in data[pid]: + if '_percent' in column: + colinfo[column] = 4 + continue + + l = len(str(data[pid][column])) + if not column in colinfo: + colinfo[column] = l + else: + if colinfo[column] < l: + colinfo[column] = l + + try: + info = ['exe', 'cmdline'] + hide = [ + int(x) if x.isdigit() else x for x in args.hide + ] + show = [ + int(x) if x.isdigit() else x for x in args.show + ] + + if not args.all and self.client.is_linux(): + hide.append(2) + + if args.info: + info = [ 'username', 'cpu_percent', 'memory_percent' ] + info + + if args.tree: + show = show or [ root ] + + for item in show: + print_pstree( + self.stdout, item, tree, data, + width=None if args.wide else width, colinfo=colinfo, info=info, + hide=hide, first=(item == root) + ) else: - for dic in outputlist: - if 'exe' in dic and not dic['exe'] and 'name' in dic and dic['name']: - dic['exe']=dic['name'].encode('utf-8', errors='replace') - if 'username' in dic and dic['username'] is None: - dic['username']="" - self.rawlog(self.formatter.table_format(outputlist, wl=columns)) - elif self.client.is_android(): - self.log(shell_exec(self.client, "ps")) - elif self.client.is_darwin(): - self.log(shell_exec(self.client, "ps aux")) - else: - self.log(shell_exec(self.client, "ps -aux")) + data = [ x for x in data if x in args.show ] if args.show else data + print_ps( + self.stdout, data, width=None if args.wide else width, + colinfo=colinfo, info=info, hide=hide + ) + + except Exception, e: + logging.exception(e) diff --git a/pupy/packages/all/pupyps.py b/pupy/packages/all/pupyps.py new file mode 100644 index 00000000..3a931b19 --- /dev/null +++ b/pupy/packages/all/pupyps.py @@ -0,0 +1,55 @@ +# -*- coding: utf-8 -*- + +import psutil + +import collections +import sys + +def pstree(): + data = {} + tree = {} + me = psutil.Process() + try: + my_user = me.username() + except: + my_user = None + + for p in psutil.process_iter(): + if not psutil.pid_exists(p.pid): + continue + + data[p.pid] = p.as_dict([ + 'name', 'username', 'cmdline', 'exe', + 'cpu_percent', 'memory_percent', 'connections' + ]) + + if p.pid == me.pid: + data[p.pid]['self'] = True + elif my_user and data[p.pid].get('username') == my_user: + data[p.pid]['same_user'] = True + + if 'connections' in data[p.pid]: + data[p.pid]['connections'] = bool(data[p.pid]['connections']) + + try: + parent = p.parent() + ppid = parent.pid if parent else 0 + if not ppid in tree: + tree[ppid] = [p.pid] + else: + tree[ppid].append(p.pid) + + except (psutil.ZombieProcess): + data[p.pid]['name'] = '< Z: ' + data[p.pid]['name'] + ' >' + + except (psutil.NoSuchProcess): + pass + + # on systems supporting PID 0, PID 0's parent is usually 0 + if 0 in tree and 0 in tree[0]: + tree[0].remove(0) + + return min(tree), tree, data + +if __name__ == '__main__': + print pstree()