Serve static image files

server/https_server.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 
 import argparse
+import os
 from gevent import pywsgi
 import ssl
 
@@ -11,6 +12,11 @@ parser.add_argument(
     dest='ca_cert',
     action='store',
     required=True)
+parser.add_argument(
+    '--image-path',
+    dest='image_path',
+    action='store',
+    required=True)
 parser.add_argument(
     '--listen-host',
     dest='listen_host',
@@ -37,17 +43,57 @@ FLAGS = parser.parse_args()
 
 class ImageRequestHandler(object):
 
+  _MIME_TYPES = {
+    '.iso': 'application/octet-stream',
+    '.json': 'application/json',
+  }
+  _BLOCK_SIZE = 2 ** 16
+
+  def __init__(self, image_path):
+    self._image_path = image_path
+
   def __call__(self, env, start_response):
-    print(env['PATH_INFO'])
-    start_response('200 OK', [('Content-Type', 'text/plain')])
-    return [b'foo']
+    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)
+
+    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('.')
+
+    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 ImageServer(object):
 
-  def __init__(self, listen_host, listen_port, server_key, server_cert, ca_cert):
+  def __init__(self, listen_host, listen_port, server_key, server_cert, ca_cert, image_path):
 
-    self._handler = ImageRequestHandler()
+    self._handler = ImageRequestHandler(image_path)
 
     self._httpd = pywsgi.WSGIServer(
         (listen_host, listen_port),
@@ -68,7 +114,8 @@ def main():
       FLAGS.listen_port,
       FLAGS.server_key,
       FLAGS.server_cert,
-      FLAGS.ca_cert)
+      FLAGS.ca_cert,
+      FLAGS.image_path)
   server.Serve()