1#!/usr/bin/env python
2
3import os
4import re
5import sys
6
7def _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
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 = None
27    return m
28
29def 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
37def 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
54class 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
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 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
123def 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
131import signal
132kSignals = {}
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 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
275if __name__ == '__main__':
276    main()
277