mxvmem

#! /usr/bin/env python3

'''
   show memory consumption per user

   The script reads conserved or live information from the /proc
   directory and shows memory consumption per user.

   The default is to read the most recent forensics dump found in
   /var/log. When called with the '-p' option, data is collected
   from /proc.

   Setting the FORENSICS_STORE environment variable overrides the
   default directory where the forensics logs are searched.


   What to do if too much memory is burned on a given host?

      #> ps -u $USER -o pid,state,lstart,cmd

   This helps finding old processes, or one may just end _all_
   own sessions on this host:

      #> killall -u $USER


   History:
   13.10.2022, created, Kreitler

'''

import os
import re
import pwd
import sys
import time
import argparse

# # Debug helper
# try:
#   from dumper import dump
# except ModuleNotFoundError:
#   from pprint import pprint as dump

# ------------------------------------------ group/summarize data in a tree

class Classifier():
  ''' Use to group/summarize values into a tree structure. '''
  def __init__(self):
    self.cnt = 0
    self.fld = dict()
    self.val = 0
    self._ik = None

  def __iter__(self): # this is either expensive, stupid, or already present elsewhere in python
    self._ik = iter(sorted (self.fld.keys()))
    return self

  def __next__(self):
    n = next(self._ik)
    return (n, self.fld[n])

  def cfy(self, fields, val=0):
    self.cnt += 1
    self.val += val
    if len(fields) > 0:
      f = fields.pop(0)
      if f not in self.fld: self.fld[f] = Classifier()
      self.fld[f].cfy(fields, val)
    return val

  def get(self, key):
    return self.fld[key]

  def getval(self):
    return self.val

  def dump_simple(self, lbl='.', ind=0):
    print('# %s %-8s %6d ,'%(' '*ind,lbl+':', self.cnt), self.val)
    ks = sorted(self.fld.keys())
    for k in ks:
      self.fld[k].dump_simple(k, ind+2)
    pass


# ------------------------------------ access to proc data (stored or live)

class ProcBuffer():
  '''
    Reads from /proc, provides an iterator. Keeps ProcStream simple,
    allows handling /proc the same way as a 'forensics' log.

    Example usage:
      pb = ProcBuffer()
      print(pb.read())
      for line in pb:
        print(line)
  '''
  def __init__(self, pdir='/proc'):
    self.pdir  = pdir
    self.buf   = []
    self.bufit = iter(self.buf)

  def __iter__(self):
    return self

  def __next__(self):
    return next(self.bufit)

  def readfile(self, fn):
    lines = 0
    try:
      f = open(fn, errors='backslashreplace')
    except (FileNotFoundError, PermissionError): # do nothing
      return 0
    for line in f:
      self.buf.append(fn + ' ' + line.strip())
      lines += 1
    f.close()
    return lines

  def read(self):
    # get machine info, return number of items
    for pf in ('uptime', 'meminfo'):
      self.readfile(self.pdir + '/' + pf)
    # go for the PIDs
    for pd in os.scandir(self.pdir):
      if pd.name.isnumeric() and pd.is_dir(): # isdigit(), isdecimal() -- which one ?
        for pf in ('stat', 'status'):
          self.readfile(self.pdir + '/' + pd.name + '/' + pf)
    return len(self.buf)

  def get_iter(self):
    self.bufit = iter(self.buf)
    return self.bufit

class ProcInfoBase():
  def __iter__(self):
    return self
  # today nobody closes files, but I like to ...
  def close(self):
    pass

class ProcInfo(ProcInfoBase):
  ''' Read from /proc .'''
  def __init__(self, pdir='/proc'):
    pb = ProcBuffer(pdir)
    pb.read()
    self.pbi = pb.get_iter()
    self.source = pdir

  def __next__(self):
    return next(self.pbi)

class ProcInfoSaved(ProcInfoBase): # aka forensics
  ''' Read from forensics file. '''
  def __init__(self, logfile):
    self.file = open(logfile, errors='backslashreplace')
    self.source = logfile

  def __next__(self):
    line = self.file.readline()
    if not line:
      self.close()
      raise StopIteration
    return line.strip()

  def close(self):
    if self.file:
      self.file.close()
      self.file = None


# ------------------------------------------------------- collect proc data

class ProcStreamParser():
  '''
      state: R is running, S is sleeping,
             D is sleeping in an uninterruptible wait,
             Z is zombie, T is traced or stopped,
             I is idle, and not in the docs ?
  '''
  def __init__(self,procfshandler):
    self.pfh      = procfshandler
    self.pidrex   = re.compile('/proc/(\d+)/(\S+)\s+(.*)') # :xxx: hardcoded '/proc'

  def parse(self, line):
    m = self.pidrex.match(line)
    if not m:
      if   line.startswith('/proc/uptime '):
        wall_clock = float(line.split()[1])
        self.pfh.set_uptime(wall_clock)
        self.pfh.report_append( '# uptime: %.2f s, %.2f d' % (wall_clock, wall_clock/3600.0/24.0) )
      elif line.startswith('/proc/meminfo MemTotal'):
        memtotal = int(line.split()[2])
        self.pfh.set_memtotal(memtotal)
        self.pfh.report_append( '# memtotal: %d kB, %.2f GB' % (memtotal, memtotal/1024/1024) )
    else:
      pid   = int(m.group(1))
      pfile = m.group(2)
      entry = m.group(3)

      if   pfile == 'stat':
        fields = entry.split()
        state = fields[2]
        start_time = float(fields[21])/100.0 # seconds please ...
        self.pfh.set_info(pid, 'start_time', start_time)
      elif pfile == 'status':
        if   entry.startswith('VmData:'):
          val = int(entry.split()[1])
          self.pfh.set_info(pid, 'vmdata', val)
        elif entry.startswith('State:'): # redundant
          val = entry.split()[1]
          self.pfh.set_info(pid, 'state', val)
        elif entry.startswith('Uid:'):
          val = int(entry.split()[1])
          self.pfh.set_info(pid, 'uid', val)


# --------------------------------------------------------------- workhorse

class ProcFsHandler():
  def __init__(self, classifier = Classifier()):
    self.Cfy      = classifier
    self.store    = {}
    self.usermap  = {}
    self.report   = []
    self.uptime   = -1
    self.memtotal = -1
    self.age_threshold = 2 * 60*60*24 # Two days

  def set_uptime(self, t):
    self.uptime = t

  def set_memtotal(self, m):
    self.memtotal = m

  def report_append(self, s):
    self.report.append(s)

  # fill hash
  def set_info(self, pid, key, val):
    if pid not in self.store:
      self.store[pid] = {}
    self.store[pid][key] = val

  def analyze(self):
    for p in self.store:
      # --------- resolve_user_names
      name = '_unknown_'
      vmdata = 0
      state = ''
      uid = self.store[p]['uid']

      if not uid in self.usermap:
        try:
          name = pwd.getpwuid(uid).pw_name
        except KeyError:
          name = str(uid)
        self.usermap[uid] = name
      name = self.usermap[uid]

      # ------------------- classify
      state = self.store[p]['state']
      if state in 'DZT': state = 'DZT'

      if 'vmdata' in self.store[p]:
        vmdata = int(self.store[p]['vmdata'])
        if (self.uptime - self.store[p]['start_time']) < self.age_threshold:
          # 'young' ones get a lower case letter
          state = state.lower()
      else:
        continue # nothing to sum up

      self.Cfy.cfy( [name, state], vmdata)


# -------------------------- memory classifier that 'reports' formated bars

class ProcMemClassifier(Classifier):

  def maxMem(self):
    m = max(k[1].val for k in self)
    return m

  def barStrings(self, scmax = 100, width = 80):
    '''
    returns a string array for display:
      string has 3 components:  'user  |X-----|  memory' -- name bar amount
    '''
    # 10 for the user, 12 is default width for the number,
    #  2 for fillers and such, 2 for the border
    bar_width = width - 10 - 12 - 2 - 2
    keys = ('DZT','S','R','dzt','s','r')
    ret = list()

    for k in self:
      user = k[0]
      mtot = k[1].val
      bar  = ' %-10s' % str(user)[:9] + '|'
      nused = 0
      for key in keys:
        if key in k[1].fld:
          ks = k[1].fld[key]
          n = int(ks.val/scmax * bar_width)
          if n > 0:
            # todo: mind max clipping!
            bar += key[0] + '-' * (n-1)
            nused += n

      remain = bar_width - nused
      if remain <= 0:
        bar += '|'
      else:
        bar += '|' + ' ' * (remain-1) + '.'
      bar += sep_fmt12(k[1].val)
      ret.append((user, mtot, bar))

    return ret

# ------------------------------------------------------------------- tools

def register_logs(logdir):
  ''' Returns forensics-logs ordered by age, newest first. '''
  tmp_logrec = []
  for log in '00 10 20 30 40 50'.split():
    logfile = os.path.join(logdir, 'forensics-%sth_min.log' % log)
    if os.access(logfile, os.R_OK):
      mt = os.stat(logfile).st_mtime
      tmp_logrec.append((logfile,mt))
  logs = sorted (tmp_logrec, key=lambda L: L[1], reverse=True)
  return logs

def sep_fmt12(n = 0.0):
  ''' Gives readable positive numbers of width 12. '''
  ret = '%12s' % '{0:,}'.format(n) # new in C: printf("%'.2f", 1234567.89); ???
  ret = ret.replace(',','.')
  if len(ret) > 12: ret = '%12.4g' % n
  return ret

def get_term_size():
  ''' Say 'no' to shutil. '''
  cols_lines = (80, 24)
  try:
    cols_lines = os.get_terminal_size()
  except (ValueError, OSError):
    pass
  return cols_lines

def chk_term_size_usable():
  ''' Make sure terminal area is large enough, tell user. '''
  bail  = False
  min_w = 40
  min_h = 12
  cols, lines = get_term_size()
  if cols < min_w:
    print('# Fatal: Terminal window is not wide enough. Columns needed:', min_w,
             '( got', cols, ')', file=sys.stderr )
    bail = True
  if lines < min_h:
    print('# Fatal: Terminal window is not high enough. Lines needed:', min_h,
             '( got', lines, ')', file=sys.stderr )
    bail = True
  if bail:
    print('# Will end now ...', file=sys.stderr)
    return False
  return True

# ---------------------------------------------------------- user interface

def handle_args():
  ''' Guess? '''
  ap = argparse.ArgumentParser(allow_abbrev=True, description =
       'Read conserved or live information from the /proc directory, ' +
       'and shows memory consumption per user.' )
  ap.add_argument('-H', dest='pydoc',
     help="show documentation", action='store_true')
  ap.add_argument("-a", dest='allentries',
     help='print all entries, makes you scroll, helps when piping', action='store_true')
  ap.add_argument("-d", dest='logdir', metavar='dir', default=None,
     help='location of forensics logs (/var/log)')
  ap.add_argument('-p', dest='readproc',
     help='read current data from proc', action='store_true')
  ap.add_argument('forensicsfile', metavar='file', nargs='?',
     help='forensics file (defaults to most recent log found)')
  return ap.parse_args()

# ------------------------------------------------------------- ab die post

if __name__ == '__main__':

  args = handle_args()

  if args.pydoc:
    import subprocess
    subprocess.run(('pydoc3', sys.argv[0]))
    quit()

  if not chk_term_size_usable(): quit()

  proc = None

  if args.forensicsfile:
    if os.access(args.forensicsfile, os.R_OK):
      proc = ProcInfoSaved(args.forensicsfile)
    else:
      print('# Error: can not read', args.forensicsfile, file=sys.stderr)
      quit()
  elif args.readproc:
    proc = ProcInfo()
  else:
    default_logdir = '/var/log'
    logdir = None
    if 'FORENSICS_STORE' in os.environ:
      logdir = os.environ['FORENSICS_STORE']
    if args.logdir:
      logdir = args.logdir
    if not logdir:
      logdir = default_logdir
    if not os.access(logdir, os.R_OK):
      print('# Error: can not access', logdir, file=sys.stderr)
      quit()
    logs = register_logs(logdir)
    if not len(logs):
      print('# Fatal: no logs found in', logdir, file=sys.stderr)
      quit()
    proc = ProcInfoSaved(logs[0][0])

  print(' Reading:', "'%s'" % proc.source, '...')

  pmc = ProcMemClassifier()
  whs = ProcFsHandler(pmc)
  psp = ProcStreamParser(whs)

  for line in proc:
    psp.parse(line)
  proc.close()

  whs.analyze()

  print('  Memory: %.1f Gb available, %.1f Gb in use (%.1f %%)\n' %
         ( whs.memtotal/1024**2,
           pmc.getval()/1024**2,
           100 * pmc.getval()/whs.memtotal )
       )

  maxmem_used = pmc.maxMem()
  cols, lines = get_term_size()

  print(' USER    ', '*** OLD/young processes ***'.center(cols-23) , 'Amount [kb]')

  mem_bars = pmc.barStrings(maxmem_used, cols)
  limit = lines-11 if not args.allentries else len(mem_bars)
  for i in sorted(mem_bars, key=lambda L: L[1], reverse=True)[:limit]:
    print(i[2])

  print('*** R = running, S = sleeping, D = deep sleep, zombie, or debugged ***'.center(cols))

  print()