iconograph/server/server.py

#!/usr/bin/python3

import argparse
import json
import os
import pyinotify
import ssl
import subprocess
import sys
import threading
import time
import uuid
from ws4py import websocket
from ws4py.server import geventserver
from ws4py.server import wsgiutils


parser = argparse.ArgumentParser(description='iconograph https_server')
parser.add_argument(
    '--ca-cert',
    dest='ca_cert',
    action='store',
    required=True)
parser.add_argument(
    '--image-path',
    dest='image_path',
    action='store',
    required=True)
parser.add_argument(
    '--image-type',
    dest='image_types',
    action='append',
    required=True)
parser.add_argument(
    '--listen-host',
    dest='listen_host',
    action='store',
    default='::')
parser.add_argument(
    '--listen-port',
    dest='listen_port',
    type=int,
    action='store',
    default=443)
parser.add_argument(
    '--server-key',
    dest='server_key',
    action='store',
    required=True)
parser.add_argument(
    '--server-cert',
    dest='server_cert',
    action='store',
    required=True)
parser.add_argument(
    '--static-path',
    dest='static_paths',
    action='append')
parser.add_argument(
    '--exec-handler',
    dest='exec_handlers',
    action='append')
FLAGS = parser.parse_args()


class WebSockets(object):
  def __init__(self):
    self.slaves = set()
    self.masters = set()
    self.targets = {}

  def __iter__(self):
    return iter(self.slaves | self.masters)

  @staticmethod
  def Broadcast(targets, msg):
    msgstr = json.dumps(msg)
    for target in targets:
      target.send(msgstr)

  def BroadcastTargets(self):
    self.Broadcast(self.masters, {
      'type': 'targets',
      'data': {
        'targets': list(self.targets.keys()),
      },
    })


class BaseWSHandler(websocket.WebSocket):
  def opened(self, image_types):
    self.send(json.dumps({
      'type': 'image_types',
      'data': {
        'image_types': list(image_types),
      },
    }))


def GetSlaveWSHandler(image_types, websockets):
  _hostname = None

  class SlaveWSHandler(BaseWSHandler):
    def opened(self):
      super().opened(image_types)
      websockets.slaves.add(self)

    def closed(self, code, reason=None):
      websockets.slaves.remove(self)
      if self._hostname:
        del websockets.targets[self._hostname]
        websockets.BroadcastTargets()

    def received_message(self, msg):
      parsed = json.loads(str(msg))
      if parsed['type'] == 'report':
        newmsg = {
          'type': 'report',
          'id': str(uuid.uuid4()),
          'received': int(time.time()),
          'client': self.peer_address,
          'data': parsed['data'],
        }
        websockets.Broadcast(websockets.masters, newmsg)

        if 'hostname' in parsed['data']:
          self._hostname = parsed['data']['hostname']
          websockets.targets[self._hostname] = self
          websockets.BroadcastTargets()

  return SlaveWSHandler


def GetMasterWSHandler(image_types, websockets):
  class MasterWSHandler(BaseWSHandler):
    def opened(self):
      super().opened(image_types)
      websockets.masters.add(self)
      self.send(json.dumps({
        'type': 'targets',
        'data': {
          'targets': list(websockets.targets.keys()),
        },
      }))

    def closed(self, code, reason=None):
      websockets.masters.remove(self)

    def received_message(self, msg):
      parsed = json.loads(str(msg))
      if parsed['type'] == 'command':
        target = parsed['target']
        if target not in websockets.targets:
          return
        newmsg = {
          'type': 'command',
          'target': target,
          'id': str(uuid.uuid4()),
          'received': int(time.time()),
          'client': self.peer_address,
          'data': parsed['data'],
        }
        websockets.targets[target].send(json.dumps(newmsg))

  return MasterWSHandler


class INotifyHandler(pyinotify.ProcessEvent):
  def __init__(self, websockets):
    self._websockets = websockets

  def process_IN_MOVED_TO(self, event):
    if event.name != 'manifest.json':
      return
    image_type = os.path.basename(event.path)
    self._websockets.Broadcast(self._websockets, {
      'type': 'new_manifest',
      'data': {
        'image_type': image_type,
      },
    })


class HTTPRequestHandler(object):

  _MIME_TYPES = {
    '.css': 'text/css',
    '.html': 'text/html',
    '.iso': 'application/octet-stream',
    '.js': 'application/javascript',
    '.json': 'application/json',
    '.tar': 'application/x-tar',
    '.woff': 'application/font-woff',
  }
  _BLOCK_SIZE = 2 ** 16

  def __init__(self, image_path, image_types, exec_handlers, static_paths, websockets):
    self._image_path = image_path
    self._image_types = image_types
    self._exec_handlers = exec_handlers
    self._static_paths = static_paths
    self._static_paths['static'] = os.path.join(os.path.dirname(sys.argv[0]), 'static')

    slave_ws_handler = GetSlaveWSHandler(image_types, websockets)
    self._slave_ws_handler = wsgiutils.WebSocketWSGIApplication(
      protocols=['iconograph-slave'],
      handler_cls=slave_ws_handler)

    master_ws_handler = GetMasterWSHandler(image_types, websockets)
    self._master_ws_handler = wsgiutils.WebSocketWSGIApplication(
      protocols=['iconograph-master'],
      handler_cls=master_ws_handler)

  def __call__(self, env, start_response):
    path = env['PATH_INFO']

    if path == '/':
      path = '/static/root.html'

    for url, file_path in self._static_paths.items():
      if path.startswith('/%s/' % url):
        file_name = path[len(url) + 2:]
        return self._ServeStaticFile(start_response, file_path, file_name)

    if path.startswith('/image/'):
      image_type, image_name = path[7:].split('/', 1)
      return self._ServeImageFile(start_response, image_type, image_name)
    elif path.startswith('/exec/'):
      method = path[6:]
      return self._ServeExec(start_response, method, env['QUERY_STRING'])
    elif path == '/ws/slave':
      return self._slave_ws_handler(env, start_response)
    elif path == '/ws/master':
      return self._master_ws_handler(env, start_response)

    start_response('404 Not found', [('Content-Type', 'text/plain')])
    return [b'Not found']

  def _MIMEType(self, file_name):
    for suffix, mime_type in self._MIME_TYPES.items():
      if file_name.endswith(suffix):
        return mime_type

  def _ServeImageFile(self, start_response, image_type, image_name):
    # Sanitize inputs
    image_type = os.path.basename(image_type)
    image_name = os.path.basename(image_name)
    assert not image_type.startswith('.')
    assert not image_name.startswith('.')
    assert image_type in self._image_types

    file_path = os.path.join(self._image_path, image_type, image_name)
    try:
      with open(file_path, 'rb') as fh:
        start_response('200 OK', [('Content-Type', self._MIMEType(image_name))])
        while True:
          block = fh.read(self._BLOCK_SIZE)
          if len(block) == 0:
            break
          yield block
    except FileNotFoundError:
      start_response('404 Not found', [('Content-Type', 'text/plain')])
      return []

  def _ServeStaticFile(self, start_response, file_path, file_name):
    file_name = os.path.basename(file_name)
    assert not file_name.startswith('.')

    full_path = os.path.join(file_path, file_name)
    try:
      with open(full_path, 'rb') as fh:
        start_response('200 OK', [('Content-Type', self._MIMEType(file_name))])
        return [fh.read()]
    except FileNotFoundError:
      start_response('404 Not found', [('Content-Type', 'text/plain')])
      return []

  def _ServeExec(self, start_response, method, arg):
    handler = self._exec_handlers[method]
    start_response('200 OK', [])
    proc = subprocess.Popen(
        [handler, arg],
        stdout=subprocess.PIPE)
    while True:
      block = proc.stdout.read(self._BLOCK_SIZE)
      if len(block) == 0:
        break
      yield block
    proc.wait()


class Server(object):

  def __init__(self, listen_host, listen_port, server_key, server_cert, ca_cert, image_path, image_types, exec_handlers, static_paths):
    websockets = WebSockets()

    wm = pyinotify.WatchManager()
    inotify_handler = INotifyHandler(websockets)
    self._notifier = pyinotify.Notifier(wm, inotify_handler)
    for image_type in image_types:
      type_path = os.path.join(image_path, image_type)
      wm.add_watch(type_path, pyinotify.IN_MOVED_TO)

    exec_handlers = dict(x.split('=', 1) for x in (exec_handlers or []))
    static_paths = dict(x.split('=', 1) for x in (static_paths or []))
    http_handler = HTTPRequestHandler(image_path, image_types, exec_handlers, static_paths, websockets)
    self._httpd = geventserver.WSGIServer(
        (listen_host, listen_port),
        http_handler,
        keyfile=server_key,
        certfile=server_cert,
        ca_certs=ca_cert,
        cert_reqs=ssl.CERT_REQUIRED,
        ssl_version=ssl.PROTOCOL_TLSv1_2)

  def Serve(self):
    self._notify_thread = threading.Thread(target=self._notifier.loop)
    self._notify_thread.daemon = True
    self._notify_thread.start()
    self._httpd.serve_forever()


def main():
  server = Server(
      FLAGS.listen_host,
      FLAGS.listen_port,
      FLAGS.server_key,
      FLAGS.server_cert,
      FLAGS.ca_cert,
      FLAGS.image_path,
      set(FLAGS.image_types),
      FLAGS.exec_handlers,
      FLAGS.static_paths)
  server.Serve()


if __name__ == '__main__':
  main()