Integrate fetcher into client.py

client/client.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 
 import argparse
+import fetcher
 import json
 import os
 import socket
@@ -14,6 +15,11 @@ parser.add_argument(
     dest='config',
     action='store',
     default='/etc/iconograph.json')
+parser.add_argument(
+    '--ca-cert',
+    dest='ca_cert',
+    action='store',
+    required=True)
 parser.add_argument(
     '--https-ca-cert',
     dest='https_ca_cert',
@@ -29,6 +35,11 @@ parser.add_argument(
     dest='https_client_key',
     action='store',
     required=True)
+parser.add_argument(
+    '--image-dir',
+    dest='image_dir',
+    action='store',
+    required=True)
 parser.add_argument(
     '--server',
     dest='server',
@@ -68,10 +79,28 @@ class Client(threadedclient.WebSocketClient):
     next_image = os.path.basename(os.readlink('/isodevice/iconograph/current'))
     return int(next_image.split('.', 1)[0])
 
+  def _OnImageTypes(self, data):
+    assert self._config['image_type'] in data['image_types']
+
+  def _OnNewManifest(self, data):
+    if data['image_type'] != self._config['image_type']:
+      return
+    fetch = fetcher.Fetcher(
+        'https://%s/image/%s' % (FLAGS.server, self._config['image_type']),
+        FLAGS.ca_cert,
+        FLAGS.image_dir,
+        FLAGS.https_ca_cert,
+        FLAGS.https_client_cert,
+        FLAGS.https_client_key)
+    fetcher.Fetch()
+    fetcher.DeleteOldImages()
+
   def received_message(self, msg):
     parsed = json.loads(msg.data.decode('utf8'))
     if parsed['type'] == 'image_types':
-      assert self._config['image_type'] in parsed['data']['image_types']
+      self._OnImageTypes(parsed['data'])
+    elif parsed['type'] == 'new_manifest':
+      self._OnNewManifest(parsed['data'])
 
 
 def main():

client/fetcher.py

@@ -1,6 +1,5 @@
 #!/usr/bin/python3
 
-import argparse
 import codecs
 import json
 import hashlib
@@ -15,43 +14,6 @@ import tempfile
 from OpenSSL import crypto
 
 
-parser = argparse.ArgumentParser(description='iconograph fetcher')
-parser.add_argument(
-    '--base-url',
-    dest='base_url',
-    action='store',
-    required=True)
-parser.add_argument(
-    '--ca-cert',
-    dest='ca_cert',
-    action='store',
-    required=True)
-parser.add_argument(
-    '--https-ca-cert',
-    dest='https_ca_cert',
-    action='store')
-parser.add_argument(
-    '--https-client-cert',
-    dest='https_client_cert',
-    action='store')
-parser.add_argument(
-    '--https-client-key',
-    dest='https_client_key',
-    action='store')
-parser.add_argument(
-    '--image-dir',
-    dest='image_dir',
-    action='store',
-    required=True)
-parser.add_argument(
-    '--max-images',
-    dest='max_images',
-    action='store',
-    type=int,
-    default=5)
-FLAGS = parser.parse_args()
-
-
 class Error(Exception):
   pass
 
@@ -204,7 +166,7 @@ class Fetcher(object):
     self._FetchImage(image)
     self._SetCurrent(image)
 
-  def DeleteOldImages(self, max_images):
+  def DeleteOldImages(self, max_images=5):
     if not max_images:
       return
     images = []
@@ -218,19 +180,3 @@ class Fetcher(object):
       print('Deleting old image:', filename)
       path = os.path.join(self._image_dir, filename)
       os.unlink(path)
-
-
-def main():
-  fetcher = Fetcher(
-      FLAGS.base_url,
-      FLAGS.ca_cert,
-      FLAGS.image_dir,
-      FLAGS.https_ca_cert,
-      FLAGS.https_client_cert,
-      FLAGS.https_client_key)
-  fetcher.Fetch()
-  fetcher.DeleteOldImages(FLAGS.max_images)
-
-
-if __name__ == '__main__':
-  main()

client/run

@@ -6,10 +6,10 @@ exec 2>&1
 
 BASE=$(dirname $0)
 
-FLAGS="$(cat /icon/config/client.flags)"
-
-HTTPS_CLIENT_KEY="/systemid/$(hostname).www.key.pem"
-HTTPS_CLIENT_CERT="/systemid/$(hostname).www.cert.pem"
-HTTPS_CA_CERT="/icon/config/ca.www.cert.pem"
-
-exec "${BASE}/client.py" --https-ca-cert="${HTTPS_CA_CERT}" --https-client-cert="${HTTPS_CLIENT_CERT}" --https-client-key="${HTTPS_CLIENT_KEY}" ${FLAGS}
+exec "${BASE}/client.py" \
+  --image-dir="/isodevice/iconograph" \
+  --ca-cert="/icon/config/ca.image.cert.pem" \
+  --https-ca-cert="/icon/config/ca.www.cert.pem" \
+  --https-client-cert="/systemid/$(hostname).www.cert.pem" \
+  --https-client-key="/systemid/$(hostname).www.key.pem" \
+  $(cat /icon/config/client.flags)