Home | History | Annotate | Download | only in Misc
      1 #!/usr/bin/env python
      2 
      3 import os
      4 import re
      5 import sys
      6 
      7 def _write_message(kind, message):
      8     import inspect, os, sys
      9 
     10     # Get the file/line where this message was generated.
     11     f = inspect.currentframe()
     12     # Step out of _write_message, and then out of wrapper.
     13     f = f.f_back.f_back
     14     file,line,_,_,_ = inspect.getframeinfo(f)
     15     location = '%s:%d' % (os.path.basename(file), line)
     16 
     17     print >>sys.stderr, '%s: %s: %s' % (location, kind, message)
     18 
     19 note = lambda message: _write_message('note', message)
     20 warning = lambda message: _write_message('warning', message)
     21 error = lambda message: (_write_message('error', message), sys.exit(1))
     22 
     23 def re_full_match(pattern, str):
     24     m = re.match(pattern, str)
     25     if m and m.end() != len(str):
     26         m = None
     27     return m
     28 
     29 def parse_time(value):
     30     minutes,value = value.split(':',1)
     31     if '.' in value:
     32         seconds,fseconds = value.split('.',1)
     33     else:
     34         seconds = value
     35     return int(minutes) * 60 + int(seconds) + float('.'+fseconds)
     36 
     37 def extractExecutable(command):
     38     """extractExecutable - Given a string representing a command line, attempt
     39     to extract the executable path, even if it includes spaces."""
     40 
     41     # Split into potential arguments.
     42     args = command.split(' ')
     43 
     44     # Scanning from the beginning, try to see if the first N args, when joined,
     45     # exist. If so that's probably the executable.
     46     for i in range(1,len(args)):
     47         cmd = ' '.join(args[:i])
     48         if os.path.exists(cmd):
     49             return cmd
     50 
     51     # Otherwise give up and return the first "argument".
     52     return args[0]
     53 
     54 class Struct:
     55     def __init__(self, **kwargs):
     56         self.fields = kwargs.keys()
     57         self.__dict__.update(kwargs)
     58 
     59     def __repr__(self):
     60         return 'Struct(%s)' % ', '.join(['%s=%r' % (k,getattr(self,k))
     61                                          for k in self.fields])
     62 
     63 kExpectedPSFields = [('PID', int, 'pid'),
     64                      ('USER', str, 'user'),
     65                      ('COMMAND', str, 'command'),
     66                      ('%CPU', float, 'cpu_percent'),
     67                      ('TIME', parse_time, 'cpu_time'),
     68                      ('VSZ', int, 'vmem_size'),
     69                      ('RSS', int, 'rss')]
     70 def getProcessTable():
     71     import subprocess
     72     p = subprocess.Popen(['ps', 'aux'], stdout=subprocess.PIPE,
     73                          stderr=subprocess.PIPE)
     74     out,err = p.communicate()
     75     res = p.wait()
     76     if p.wait():
     77         error('unable to get process table')
     78     elif err.strip():
     79         error('unable to get process table: %s' % err)
     80 
     81     lns = out.split('\n')
     82     it = iter(lns)
     83     header = it.next().split()
     84     numRows = len(header)
     85 
     86     # Make sure we have the expected fields.
     87     indexes = []
     88     for field in kExpectedPSFields:
     89         try:
     90             indexes.append(header.index(field[0]))
     91         except:
     92             if opts.debug:
     93                 raise
     94             error('unable to get process table, no %r field.' % field[0])
     95 
     96     table = []
     97     for i,ln in enumerate(it):
     98         if not ln.strip():
     99             continue
    100 
    101         fields = ln.split(None, numRows - 1)
    102         if len(fields) != numRows:
    103             warning('unable to process row: %r' % ln)
    104             continue
    105 
    106         record = {}
    107         for field,idx in zip(kExpectedPSFields, indexes):
    108             value = fields[idx]
    109             try:
    110                 record[field[2]] = field[1](value)
    111             except:
    112                 if opts.debug:
    113                     raise
    114                 warning('unable to process %r in row: %r' % (field[0], ln))
    115                 break
    116         else:
    117             # Add our best guess at the executable.
    118             record['executable'] = extractExecutable(record['command'])
    119             table.append(Struct(**record))
    120 
    121     return table
    122 
    123 def getSignalValue(name):
    124     import signal
    125     if name.startswith('SIG'):
    126         value = getattr(signal, name)
    127         if value and isinstance(value, int):
    128             return value
    129     error('unknown signal: %r' % name)
    130 
    131 import signal
    132 kSignals = {}
    133 for name in dir(signal):
    134     if name.startswith('SIG') and name == name.upper() and name.isalpha():
    135         kSignals[name[3:]] = getattr(signal, name)
    136 
    137 def main():
    138     global opts
    139     from optparse import OptionParser, OptionGroup
    140     parser = OptionParser("usage: %prog [options] {pid}*")
    141 
    142     # FIXME: Add -NNN and -SIGNAME options.
    143 
    144     parser.add_option("-s", "", dest="signalName",
    145                       help="Name of the signal to use (default=%default)",
    146                       action="store", default='INT',
    147                       choices=kSignals.keys())
    148     parser.add_option("-l", "", dest="listSignals",
    149                       help="List known signal names",
    150                       action="store_true", default=False)
    151 
    152     parser.add_option("-n", "--dry-run", dest="dryRun",
    153                       help="Only print the actions that would be taken",
    154                       action="store_true", default=False)
    155     parser.add_option("-v", "--verbose", dest="verbose",
    156                       help="Print more verbose output",
    157                       action="store_true", default=False)
    158     parser.add_option("", "--debug", dest="debug",
    159                       help="Enable debugging output",
    160                       action="store_true", default=False)
    161     parser.add_option("", "--force", dest="force",
    162                       help="Perform the specified commands, even if it seems like a bad idea",
    163                       action="store_true", default=False)
    164 
    165     inf = float('inf')
    166     group = OptionGroup(parser, "Process Filters")
    167     group.add_option("", "--name", dest="execName", metavar="REGEX",
    168                       help="Kill processes whose name matches the given regexp",
    169                       action="store", default=None)
    170     group.add_option("", "--exec", dest="execPath", metavar="REGEX",
    171                       help="Kill processes whose executable matches the given regexp",
    172                       action="store", default=None)
    173     group.add_option("", "--user", dest="userName", metavar="REGEX",
    174                       help="Kill processes whose user matches the given regexp",
    175                       action="store", default=None)
    176     group.add_option("", "--min-cpu", dest="minCPU", metavar="PCT",
    177                       help="Kill processes with CPU usage >= PCT",
    178                       action="store", type=float, default=None)
    179     group.add_option("", "--max-cpu", dest="maxCPU", metavar="PCT",
    180                       help="Kill processes with CPU usage <= PCT",
    181                       action="store", type=float, default=inf)
    182     group.add_option("", "--min-mem", dest="minMem", metavar="N",
    183                       help="Kill processes with virtual size >= N (MB)",
    184                       action="store", type=float, default=None)
    185     group.add_option("", "--max-mem", dest="maxMem", metavar="N",
    186                       help="Kill processes with virtual size <= N (MB)",
    187                       action="store", type=float, default=inf)
    188     group.add_option("", "--min-rss", dest="minRSS", metavar="N",
    189                       help="Kill processes with RSS >= N",
    190                       action="store", type=float, default=None)
    191     group.add_option("", "--max-rss", dest="maxRSS", metavar="N",
    192                       help="Kill processes with RSS <= N",
    193                       action="store", type=float, default=inf)
    194     group.add_option("", "--min-time", dest="minTime", metavar="N",
    195                       help="Kill processes with CPU time >= N (seconds)",
    196                       action="store", type=float, default=None)
    197     group.add_option("", "--max-time", dest="maxTime", metavar="N",
    198                       help="Kill processes with CPU time <= N (seconds)",
    199                       action="store", type=float, default=inf)
    200     parser.add_option_group(group)
    201 
    202     (opts, args) = parser.parse_args()
    203 
    204     if opts.listSignals:
    205         items = [(v,k) for k,v in kSignals.items()]
    206         items.sort()
    207         for i in range(0, len(items), 4):
    208             print '\t'.join(['%2d) SIG%s' % (k,v)
    209                              for k,v in items[i:i+4]])
    210         sys.exit(0)
    211 
    212     # Figure out the signal to use.
    213     signal = kSignals[opts.signalName]
    214     signalValueName = str(signal)
    215     if opts.verbose:
    216         name = dict((v,k) for k,v in kSignals.items()).get(signal,None)
    217         if name:
    218             signalValueName = name
    219             note('using signal %d (SIG%s)' % (signal, name))
    220         else:
    221             note('using signal %d' % signal)
    222 
    223     # Get the pid list to consider.
    224     pids = set()
    225     for arg in args:
    226         try:
    227             pids.add(int(arg))
    228         except:
    229             parser.error('invalid positional argument: %r' % arg)
    230 
    231     filtered = ps = getProcessTable()
    232 
    233     # Apply filters.
    234     if pids:
    235         filtered = [p for p in filtered
    236                     if p.pid in pids]
    237     if opts.execName is not None:
    238         filtered = [p for p in filtered
    239                     if re_full_match(opts.execName,
    240                                      os.path.basename(p.executable))]
    241     if opts.execPath is not None:
    242         filtered = [p for p in filtered
    243                     if re_full_match(opts.execPath, p.executable)]
    244     if opts.userName is not None:
    245         filtered = [p for p in filtered
    246                     if re_full_match(opts.userName, p.user)]
    247     filtered = [p for p in filtered
    248                 if opts.minCPU <= p.cpu_percent <= opts.maxCPU]
    249     filtered = [p for p in filtered
    250                 if opts.minMem <= float(p.vmem_size) / (1<<20) <= opts.maxMem]
    251     filtered = [p for p in filtered
    252                 if opts.minRSS <= p.rss <= opts.maxRSS]
    253     filtered = [p for p in filtered
    254                 if opts.minTime <= p.cpu_time <= opts.maxTime]
    255 
    256     if len(filtered) == len(ps):
    257         if not opts.force and not opts.dryRun:
    258             error('refusing to kill all processes without --force')
    259 
    260     if not filtered:
    261         warning('no processes selected')
    262 
    263     for p in filtered:
    264         if opts.verbose:
    265             note('kill(%r, %s) # (user=%r, executable=%r, CPU=%2.2f%%, time=%r, vmem=%r, rss=%r)' %
    266                  (p.pid, signalValueName, p.user, p.executable, p.cpu_percent, p.cpu_time, p.vmem_size, p.rss))
    267         if not opts.dryRun:
    268             try:
    269                 os.kill(p.pid, signal)
    270             except OSError:
    271                 if opts.debug:
    272                     raise
    273                 warning('unable to kill PID: %r' % p.pid)
    274 
    275 if __name__ == '__main__':
    276     main()
    277