wrap.py (8704B)
1 #Embedded file name: C:\BuildAgent\work\20d73f1b24eed06e\eve\release\V24.01\packages\datadog\dogshell\wrap.py 2 from __future__ import print_function 3 import optparse 4 import subprocess 5 import sys 6 import threading 7 import time 8 from datadog import initialize, api 9 from datadog.util.config import get_version 10 SUCCESS = 'success' 11 ERROR = 'error' 12 MAX_EVENT_BODY_LENGTH = 3000 13 14 class Timeout(Exception): 15 pass 16 17 18 class OutputReader(threading.Thread): 19 20 def __init__(self, proc_out, fwd_out = None): 21 threading.Thread.__init__(self) 22 self.daemon = True 23 self._out_content = u'' 24 self._out = proc_out 25 self._fwd_out = fwd_out 26 27 def run(self): 28 for line in iter(self._out.readline, ''): 29 if self._fwd_out is not None: 30 self._fwd_out.write(line) 31 line = line.decode('utf-8') 32 self._out_content += line 33 34 self._out.close() 35 36 @property 37 def content(self): 38 return self._out_content 39 40 41 def poll_proc(proc, sleep_interval, timeout): 42 start_time = time.time() 43 returncode = None 44 while returncode is None: 45 returncode = proc.poll() 46 if time.time() - start_time > timeout: 47 raise Timeout() 48 else: 49 time.sleep(sleep_interval) 50 51 return returncode 52 53 54 def execute(cmd, cmd_timeout, sigterm_timeout, sigkill_timeout, proc_poll_interval, buffer_outs): 55 start_time = time.time() 56 returncode = -1 57 stdout = '' 58 stderr = '' 59 try: 60 proc = subprocess.Popen(u' '.join(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) 61 except Exception: 62 print(u'Failed to execute %s' % repr(cmd), file=sys.stderr) 63 raise 64 65 try: 66 out_reader = OutputReader(proc.stdout, sys.stdout if not buffer_outs else None) 67 err_reader = OutputReader(proc.stderr, sys.stderr if not buffer_outs else None) 68 out_reader.start() 69 err_reader.start() 70 returncode = poll_proc(proc, proc_poll_interval, cmd_timeout) 71 out_reader.join() 72 err_reader.join() 73 stdout = out_reader.content 74 stderr = err_reader.content 75 duration = time.time() - start_time 76 except Timeout: 77 duration = time.time() - start_time 78 try: 79 proc.terminate() 80 sigterm_start = time.time() 81 try: 82 print('Command timed out after %.2fs, killing with SIGTERM' % (time.time() - start_time), file=sys.stderr) 83 poll_proc(proc, proc_poll_interval, sigterm_timeout) 84 returncode = Timeout 85 except Timeout: 86 print('SIGTERM timeout failed after %.2fs, killing with SIGKILL' % (time.time() - sigterm_start), file=sys.stderr) 87 proc.kill() 88 poll_proc(proc, proc_poll_interval, sigkill_timeout) 89 returncode = Timeout 90 91 except OSError as e: 92 if e.errno != 3: 93 raise 94 95 return (returncode, 96 stdout, 97 stderr, 98 duration) 99 100 101 def trim_text(text, max_len): 102 if len(text) <= max_len: 103 return text 104 trimmed_text = u'{top_third}\n```\n*...trimmed...*\n```\n{bottom_two_third}\n'.format(top_third=text[:max_len / 3], bottom_two_third=text[len(text) - 2 * max_len / 3:]) 105 return trimmed_text 106 107 108 def build_event_body(cmd, returncode, stdout, stderr, notifications): 109 fmt_stdout = u'' 110 fmt_stderr = u'' 111 fmt_notifications = u'' 112 max_length = MAX_EVENT_BODY_LENGTH / 2 if stdout and stderr else MAX_EVENT_BODY_LENGTH 113 if stdout: 114 fmt_stdout = u'**>>>> STDOUT <<<<**\n```\n{stdout} \n```\n'.format(stdout=trim_text(stdout, max_length)) 115 if stderr: 116 fmt_stderr = u'**>>>> STDERR <<<<**\n```\n{stderr} \n```\n'.format(stderr=trim_text(stderr, max_length)) 117 if notifications: 118 fmt_notifications = u'**>>>> NOTIFICATIONS <<<<**\n\n {notifications}\n'.format(notifications=notifications) 119 return u'%%%\n**>>>> CMD <<<<**\n```\n{command} \n```\n**>>>> EXIT CODE <<<<**\n\n {returncode}\n\n\n{stdout}{stderr}{notifications}%%%\n'.format(command=u' '.join(cmd), returncode=returncode, stdout=fmt_stdout, stderr=fmt_stderr, notifications=fmt_notifications) 120 121 122 def main(): 123 parser = optparse.OptionParser(usage='%prog -n [event_name] -k [api_key] --submit_mode [ all | errors ] [options] "command". \n\nNote that you need to enclose your command in quotes to prevent python as soon as there is a space in your command. \n \nNOTICE: In normal mode, the whole stderr is printed before stdout, in flush_live mode they will be mixed but there is not guarantee that messages sent by the command on both stderr and stdout are printed in the order they were sent.', version='%prog {0}'.format(get_version())) 124 parser.add_option('-n', '--name', action='store', type='string', help='the name of the event as it should appear on your Datadog stream') 125 parser.add_option('-k', '--api_key', action='store', type='string', help='your DataDog API Key') 126 parser.add_option('-m', '--submit_mode', action='store', type='choice', default='errors', choices=['errors', 'all'], help='[ all | errors ] if set to error, an event will be sent only of the command exits with a non zero exit status or if it times out.') 127 parser.add_option('-p', '--priority', action='store', type='choice', choices=['normal', 'low'], help="the priority of the event (default: 'normal')") 128 parser.add_option('-t', '--timeout', action='store', type='int', default=86400, help='(in seconds) a timeout after which your command must be aborted. An event will be sent to your DataDog stream (default: 24hours)') 129 parser.add_option('--sigterm_timeout', action='store', type='int', default=120, help='(in seconds) When your command times out, the process it triggers is sent a SIGTERM. If this sigterm_timeout is reached, it will be sent a SIGKILL signal. (default: 2m)') 130 parser.add_option('--sigkill_timeout', action='store', type='int', default=60, help='(in seconds) how long to wait at most after SIGKILL has been sent (default: 60s)') 131 parser.add_option('--proc_poll_interval', action='store', type='float', default=0.5, help='(in seconds). interval at which your command will be polled (default: 500ms)') 132 parser.add_option('--notify_success', action='store', type='string', default='', help='a message string and @people directives to send notifications in case of success.') 133 parser.add_option('--notify_error', action='store', type='string', default='', help='a message string and @people directives to send notifications in case of error.') 134 parser.add_option('-b', '--buffer_outs', action='store_true', dest='buffer_outs', default=False, help='displays the stderr and stdout of the command only once it has returned (the command outputs remains buffered in dogwrap meanwhile)') 135 parser.add_option('--tags', action='store', type='string', dest='tags', default='', help='comma separated list of tags') 136 options, args = parser.parse_args() 137 cmd = [] 138 for part in args: 139 cmd.extend(part.split(' ')) 140 141 returncode, stdout, stderr, duration = execute(cmd, options.timeout, options.sigterm_timeout, options.sigkill_timeout, options.proc_poll_interval, options.buffer_outs) 142 initialize(api_key=options.api_key) 143 host = api._host_name 144 if returncode == 0: 145 alert_type = SUCCESS 146 event_priority = 'low' 147 event_title = u'[%s] %s succeeded in %.2fs' % (host, options.name, duration) 148 else: 149 alert_type = ERROR 150 event_priority = 'normal' 151 if returncode is Timeout: 152 event_title = u'[%s] %s timed out after %.2fs' % (host, options.name, duration) 153 returncode = -1 154 else: 155 event_title = u'[%s] %s failed in %.2fs' % (host, options.name, duration) 156 notifications = '' 157 if alert_type == SUCCESS and options.notify_success: 158 notifications = options.notify_success 159 elif alert_type == ERROR and options.notify_error: 160 notifications = options.notify_error 161 if options.tags: 162 tags = [ t.strip() for t in options.tags.split(',') ] 163 else: 164 tags = None 165 event_body = build_event_body(cmd, returncode, stdout, stderr, notifications) 166 event = {'alert_type': alert_type, 167 'aggregation_key': options.name, 168 'host': host, 169 'priority': options.priority or event_priority, 170 'tags': tags} 171 if options.buffer_outs: 172 print(stderr.strip(), file=sys.stderr) 173 print(stdout.strip(), file=sys.stdout) 174 if options.submit_mode == 'all' or returncode != 0: 175 api.Event.create(title=event_title, text=event_body, **event) 176 sys.exit(returncode) 177 178 179 if __name__ == '__main__': 180 main()