Rename to server.py

server/server.py

@@ -0,0 +1,185 @@
+#!/usr/bin/python3
+
+import argparse
+import json
+import os
+import pyinotify
+import ssl
+import sys
+import threading
+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)
+FLAGS = parser.parse_args()
+
+
+def GetWebSocketHandler(image_types, websockets):
+  class WebSocketHandler(websocket.WebSocket):
+    def opened(self):
+      self.send(json.dumps({
+        'type': 'image_types',
+        'data': {
+          'image_types': list(image_types),
+        },
+      }), False)
+      websockets.add(self)
+
+    def closed(self, code, reason=None):
+      websockets.remove(self)
+
+  return WebSocketHandler
+
+
+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)
+    for websocket in self._websockets:
+      websocket.send(json.dumps({
+        'type': 'new_manifest',
+        'data': {
+          'image_type': image_type,
+        },
+      }), False)
+
+
+class HTTPRequestHandler(object):
+
+  _MIME_TYPES = {
+    '.iso': 'application/octet-stream',
+    '.json': 'application/json',
+  }
+  _BLOCK_SIZE = 2 ** 16
+
+  def __init__(self, image_path, image_types, websockets):
+    self._image_path = image_path
+    self._image_types = image_types
+    inner_handler = GetWebSocketHandler(image_types, websockets)
+    self._websocket_handler = wsgiutils.WebSocketWSGIApplication(handler_cls=inner_handler)
+
+  def __call__(self, env, start_response):
+    path = env['PATH_INFO']
+    if path.startswith('/image/'):
+      image_type, image_name = path[7:].split('/', 1)
+      return self._ServeImageFile(start_response, image_type, image_name)
+    elif path == '/ws':
+      return self._websocket_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')
+      return []
+
+    return
+
+
+class Server(object):
+
+  def __init__(self, listen_host, listen_port, server_key, server_cert, ca_cert, image_path, image_types):
+    websockets = set()
+
+    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)
+
+    http_handler = HTTPRequestHandler(image_path, image_types, 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))
+  server.Serve()
+
+
+if __name__ == '__main__':
+  main()