From f8527fe33905f2f2c7597ac95f27493f67c44d85 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <ol.shevchuk@samsung.com>
Date: Thu, 16 Mar 2017 14:04:21 +0200
Subject: [PATCH] New (psutil based) ps module

---
 pupy/modules/ps.py          | 202 +++++++++++++++++++++++++++++++-----
 pupy/packages/all/pupyps.py |  55 ++++++++++
 2 files changed, 230 insertions(+), 27 deletions(-)
 create mode 100644 pupy/packages/all/pupyps.py

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()