services/nodemanager/arvnodeman/status.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: AGPL-3.0
   4
   5 from __future__ import absolute_import, print_function
   6 from future import standard_library
   7
   8 import http.server
   9 import json
  10 import logging
  11 import socketserver
  12 import threading
  13
  14 _logger = logging.getLogger('status.Handler')
  15
  16
  17 class Server(socketserver.ThreadingMixIn, http.server.HTTPServer, object):
  18     def __init__(self, config):
  19         port = config.getint('Manage', 'port')
  20         self.enabled = port >= 0
  21         if not self.enabled:
  22             _logger.warning("Management server disabled. "+
  23                             "Use [Manage] config section to enable.")
  24             return
  25         self._config = config
  26         self._tracker = tracker
  27         super(Server, self).__init__(
  28             (config.get('Manage', 'address'), port), Handler)
  29         self._thread = threading.Thread(target=self.serve_forever)
  30         self._thread.daemon = True
  31
  32     def start(self):
  33         if self.enabled:
  34             self._thread.start()
  35
  36
  37 class Handler(http.server.BaseHTTPRequestHandler, object):
  38     def do_GET(self):
  39         if self.path == '/status.json':
  40             self.send_response(200)
  41             self.send_header('Content-type', 'application/json')
  42             self.end_headers()
  43             self.wfile.write(tracker.get_json())
  44         elif self.path == '/_health/ping':
  45             code, msg = self.check_auth()
  46
  47             if code != 200:
  48               self.send_response(code)
  49               self.wfile.write(msg)
  50             else:
  51               self.send_response(200)
  52               self.send_header('Content-type', 'application/json')
  53               self.end_headers()
  54               self.wfile.write(json.dumps({"health":"OK"}))
  55         else:
  56             self.send_response(404)
  57
  58     def log_message(self, fmt, *args, **kwargs):
  59         _logger.info(fmt, *args, **kwargs)
  60
  61     def check_auth(self):
  62         mgmt_token = self.server._config.get('Manage', 'ManagementToken')
  63         auth_header = self.headers.get('Authorization', None)
  64
  65         if mgmt_token == '':
  66           return 404, "disabled"
  67         elif auth_header == None:
  68           return 401, "authorization required"
  69         elif auth_header != 'Bearer '+mgmt_token:
  70           return 403, "authorization error"
  71         return 200, ""
  72
  73 class Tracker(object):
  74     def __init__(self):
  75         self._mtx = threading.Lock()
  76         self._latest = {}
  77
  78     def get_json(self):
  79         with self._mtx:
  80             return json.dumps(self._latest)
  81
  82     def keys(self):
  83         with self._mtx:
  84             return self._latest.keys()
  85
  86     def update(self, updates):
  87         with self._mtx:
  88             self._latest.update(updates)
  89
  90
  91 tracker = Tracker()