Eli Friedman | 77a1fe9 | 2009-07-10 20:15:12 +0000 | [diff] [blame] | 1 | #!/usr/bin/env python |
Daniel Dunbar | 02ca6ea | 2009-03-04 21:06:52 +0000 | [diff] [blame] | 2 | |
| 3 | """CaptureCmd - A generic tool for capturing information about the |
| 4 | invocations of another program. |
| 5 | |
| 6 | Usage |
| 7 | -- |
| 8 | 1. Move the original tool to a safe known location. |
| 9 | |
| 10 | 2. Link CaptureCmd to the original tool's location. |
| 11 | |
| 12 | 3. Define CAPTURE_CMD_PROGRAM to the known location of the original |
| 13 | tool; this must be an absolute path. |
| 14 | |
| 15 | 4. Define CAPTURE_CMD_DIR to a directory to write invocation |
| 16 | information to. |
| 17 | """ |
| 18 | |
| 19 | import hashlib |
| 20 | import os |
| 21 | import sys |
| 22 | import time |
| 23 | |
| 24 | def saveCaptureData(prefix, dir, object): |
| 25 | string = repr(object) + '\n' |
| 26 | key = hashlib.sha1(string).hexdigest() |
| 27 | path = os.path.join(dir, |
| 28 | prefix + key) |
| 29 | if not os.path.exists(path): |
| 30 | f = open(path, 'wb') |
| 31 | f.write(string) |
| 32 | f.close() |
| 33 | return prefix + key |
| 34 | |
| 35 | def main(): |
| 36 | program = os.getenv('CAPTURE_CMD_PROGRAM') |
| 37 | dir = os.getenv('CAPTURE_CMD_DIR') |
| 38 | fallback = os.getenv('CAPTURE_CMD_FALLBACK') |
| 39 | if not program: |
| 40 | raise ValueError('CAPTURE_CMD_PROGRAM is not defined!') |
| 41 | if not dir: |
| 42 | raise ValueError('CAPTURE_CMD_DIR is not defined!') |
| 43 | |
| 44 | # Make the output directory if it doesn't already exist. |
| 45 | if not os.path.exists(dir): |
| 46 | os.mkdir(dir, 0700) |
| 47 | |
| 48 | # Get keys for various data. |
| 49 | env = os.environ.items() |
| 50 | env.sort() |
| 51 | envKey = saveCaptureData('env-', dir, env) |
| 52 | cwdKey = saveCaptureData('cwd-', dir, os.getcwd()) |
| 53 | argvKey = saveCaptureData('argv-', dir, sys.argv) |
| 54 | entry = (time.time(), envKey, cwdKey, argvKey) |
| 55 | saveCaptureData('cmd-', dir, entry) |
| 56 | |
| 57 | if fallback: |
| 58 | pid = os.fork() |
| 59 | if not pid: |
| 60 | os.execv(program, sys.argv) |
| 61 | os._exit(1) |
| 62 | else: |
| 63 | res = os.waitpid(pid, 0) |
| 64 | if not res: |
| 65 | os.execv(fallback, sys.argv) |
| 66 | os._exit(1) |
| 67 | os._exit(res) |
| 68 | else: |
| 69 | os.execv(program, sys.argv) |
| 70 | os._exit(1) |
| 71 | |
| 72 | if __name__ == '__main__': |
| 73 | main() |