code.grnet.gr Git - snf-image/blob - snf-image-host/helper-monitor.py

   1 #!/usr/bin/env python
   2
   3 # Copyright (C) 2012 GRNET S.A.
   4 #
   5 # This program is free software; you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation; either version 2 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful, but
  11 # WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  13 # General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with this program; if not, write to the Free Software
  17 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  18 # 02110-1301, USA.
  19
  20 import sys
  21 import os
  22 import time
  23 import json
  24 import re
  25
  26 LINESIZE = 512
  27 BUFSIZE = 512
  28 PROGNAME = os.path.basename(sys.argv[0])
  29 STDERR_MAXLINES = 10
  30 MAXLINES = 100
  31 MSG_TYPE = 'image-helper'
  32
  33 PROTOCOL = {
  34     'TASK_START': ('task-start', 'task'),
  35     'TASK_END': ('task-end', 'task'),
  36     'WARNING': ('warning', 'messages'),
  37     'STDERR': ('error', 'stderr'),
  38     'ERROR': ('error', 'messages')}
  39
  40
  41 def error(msg):
  42     sys.stderr.write("HELPER-MONITOR ERROR: %s\n" % msg)
  43     sys.exit(1)
  44
  45
  46 def send(fd, msg_type, value):
  47     subtype, value_name = PROTOCOL[msg_type]
  48
  49     msg = {}
  50     msg['type'] = MSG_TYPE
  51     msg['subtype'] = subtype
  52     msg[value_name] = value
  53     msg['timestamp'] = time.time()
  54     os.write(fd, "%s\n" % json.dumps(msg))
  55
  56
  57 if __name__ == "__main__":
  58     usage = "Usage: %s <file-descriptor>\n" % PROGNAME
  59
  60     if len(sys.argv) != 2:
  61         sys.stderr.write(usage)
  62         sys.exit(1)
  63
  64     try:
  65         fd = int(sys.argv[1])
  66     except ValueError:
  67         error("File descriptor is not an integer")
  68
  69     try:
  70         os.fstat(fd)
  71     except OSError:
  72         error("File descriptor is not valid")
  73
  74     lines_left = 0
  75     line_count = 0
  76     stderr = ""
  77     line = ""
  78     while True:
  79         # Can't use sys.stdin.readline since I want unbuffered I/O
  80         new_data = os.read(sys.stdin.fileno(), BUFSIZE)
  81
  82         if not new_data:
  83             if not line:
  84                 break
  85             else:
  86                 new_data = '\n'
  87
  88         while True:
  89             split = new_data.split('\n', 1)
  90             line += split[0]
  91             if len(split) == 1:
  92                 if len(line) > LINESIZE:
  93                     error("Line size exceeded the maximum allowed size")
  94                 break
  95
  96             new_data = split[1]
  97
  98             line_count += 1
  99             if line_count >= MAXLINES + 1:
 100                 error("Exceeded maximum allowed number of lines: %d." %
 101                       MAXLINES)
 102
 103             if lines_left > 0:
 104                 stderr += "%s\n" % line
 105                 lines_left -= 1
 106                 if lines_left == 0:
 107                     send(fd, "STDERR", stderr)
 108                     stderr = ""
 109                 line = ""
 110                 continue
 111
 112             line = line.strip()
 113             if len(line) == 0:
 114                 continue
 115
 116             if line.startswith("STDERR:"):
 117                 m = re.match("STDERR:(\d+):(.*)", line)
 118                 if not m:
 119                     error("Invalid syntax for STDERR line")
 120                 try:
 121                     lines_left = int(m.group(1))
 122                 except ValueError:
 123                     error("Second field in STDERR line must be an integer")
 124
 125                 if lines_left > STDERR_MAXLINES:
 126                     error("Too many lines in the STDERR output")
 127                 elif lines_left < 0:
 128                     error("Second field of STDERR: %d is invalid" % lines_left)
 129
 130                 if lines_left > 0:
 131                     stderr = m.group(2) + "\n"
 132                     lines_left -= 1
 133
 134                 if lines_left == 0:
 135                     send(fd, "STDERR", stderr)
 136                     stderr = ""
 137             elif line.startswith("TASK_START:") \
 138                 or line.startswith("TASK_END:") \
 139                 or line.startswith("WARNING:") \
 140                 or line.startswith("ERROR:"):
 141                 (msg_type, _, value) = line.partition(':')
 142
 143                 if line.startswith("WARNING:") or line.startswith("ERROR:"):
 144                     value = [value]
 145                 send(fd, msg_type, value)
 146             else:
 147                 error("Unknown command!")
 148
 149             # Remove the processed line
 150             line = ""
 151
 152 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :