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