277 lines · plain
1#!/usr/bin/env python2 3import os4import re5import sys6 7def _write_message(kind, message):8 import inspect, os, sys9 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_back14 file,line,_,_,_ = inspect.getframeinfo(f)15 location = '%s:%d' % (os.path.basename(file), line)16 17 print('%s: %s: %s' % (location, kind, message), file=sys.stderr)18 19note = lambda message: _write_message('note', message)20warning = lambda message: _write_message('warning', message)21error = lambda message: (_write_message('error', message), sys.exit(1))22 23def re_full_match(pattern, str):24 m = re.match(pattern, str)25 if m and m.end() != len(str):26 m = None27 return m28 29def parse_time(value):30 minutes,value = value.split(':',1)31 if '.' in value:32 seconds,fseconds = value.split('.',1)33 else:34 seconds = value35 return int(minutes) * 60 + int(seconds) + float('.'+fseconds)36 37def extractExecutable(command):38 """extractExecutable - Given a string representing a command line, attempt39 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 cmd50 51 # Otherwise give up and return the first "argument".52 return args[0]53 54class Struct:55 def __init__(self, **kwargs):56 self.fields = list(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 63kExpectedPSFields = [('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')]70def getProcessTable():71 import subprocess72 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 raise94 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 continue100 101 fields = ln.split(None, numRows - 1)102 if len(fields) != numRows:103 warning('unable to process row: %r' % ln)104 continue105 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 raise114 warning('unable to process %r in row: %r' % (field[0], ln))115 break116 else:117 # Add our best guess at the executable.118 record['executable'] = extractExecutable(record['command'])119 table.append(Struct(**record))120 121 return table122 123def getSignalValue(name):124 import signal125 if name.startswith('SIG'):126 value = getattr(signal, name)127 if value and isinstance(value, int):128 return value129 error('unknown signal: %r' % name)130 131import signal132kSignals = {}133for name in dir(signal):134 if name.startswith('SIG') and name == name.upper() and name.isalpha():135 kSignals[name[3:]] = getattr(signal, name)136 137def main():138 global opts139 from optparse import OptionParser, OptionGroup140 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=list(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 list(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 list(kSignals.items())).get(signal,None)217 if name:218 signalValueName = name219 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 filtered236 if p.pid in pids]237 if opts.execName is not None:238 filtered = [p for p in filtered239 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 filtered243 if re_full_match(opts.execPath, p.executable)]244 if opts.userName is not None:245 filtered = [p for p in filtered246 if re_full_match(opts.userName, p.user)]247 filtered = [p for p in filtered248 if opts.minCPU <= p.cpu_percent <= opts.maxCPU]249 filtered = [p for p in filtered250 if opts.minMem <= float(p.vmem_size) / (1<<20) <= opts.maxMem]251 filtered = [p for p in filtered252 if opts.minRSS <= p.rss <= opts.maxRSS]253 filtered = [p for p in filtered254 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 raise273 warning('unable to kill PID: %r' % p.pid)274 275if __name__ == '__main__':276 main()277