From 1aa49754cceadbab898f8182ed7fe4adb9d14220 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <flamingcow@flamingcow-macbookpro.roam.corp.google.com>
Date: Mon, 9 May 2016 15:39:11 -0700
Subject: [PATCH] Make client report more data, use system config

---
 client/client.py | 42 +++++++++++++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/client/client.py b/client/client.py
index 41070ce..9e71b55 100755
--- a/client/client.py
+++ b/client/client.py
@@ -2,6 +2,7 @@
 
 import argparse
 import json
+import os
 import socket
 import time
 from ws4py.client import threadedclient
@@ -9,10 +10,10 @@ from ws4py.client import threadedclient
 
 parser = argparse.ArgumentParser(description='iconograph fetcher')
 parser.add_argument(
-    '--server',
-    dest='server',
+    '--config',
+    dest='config',
     action='store',
-    required=True)
+    default='/etc/iconograph.json')
 parser.add_argument(
     '--https-ca-cert',
     dest='https_ca_cert',
@@ -29,8 +30,8 @@ parser.add_argument(
     action='store',
     required=True)
 parser.add_argument(
-    '--image-type',
-    dest='image_type',
+    '--server',
+    dest='server',
     action='store',
     required=True)
 FLAGS = parser.parse_args()
@@ -38,17 +39,24 @@ FLAGS = parser.parse_args()
 
 class Client(threadedclient.WebSocketClient):
 
+  def __init__(self, config_path, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    with open(config_path, 'r') as fh:
+      self._config = json.load(fh)
+
   def Loop(self):
     self.daemon = True
     self.connect()
     while True:
+      report = {
+        'hostname': socket.gethostname(),
+        'uptime_seconds': self._Uptime(),
+        'next_timestamp': self._NextTimestamp(),
+      }
+      report.update(self._config)
       self.send(json.dumps({
         'type': 'report',
-        'data': {
-          'hostname': socket.gethostname(),
-          'uptime_seconds': self._Uptime(),
-          'image_type': FLAGS.image_type,
-        },
+        'data': report,
       }), False)
       time.sleep(5.0)
 
@@ -56,10 +64,14 @@ class Client(threadedclient.WebSocketClient):
     with open('/proc/uptime', 'r') as fh:
       return int(float(fh.readline().split(' ', 1)[0]))
 
+  def _NextTimestamp(self):
+    next_image = os.path.basename(os.readlink('/isodevice/iconograph/current'))
+    return int(next_image.split('.', 1)[0])
+
   def received_message(self, msg):
-    parsed = json.loads(str(msg))
+    parsed = json.loads(msg.data.decode('utf8'))
     if parsed['type'] == 'image_types':
-      assert FLAGS.image_type in parsed['data']['image_types']
+      assert self._config['image_type'] in parsed['data']['image_types']
 
 
 def main():
@@ -68,7 +80,11 @@ def main():
     'certfile': FLAGS.https_client_cert,
     'ca_certs': FLAGS.https_ca_cert,
   }
-  client = Client('wss://%s/ws/slave' % FLAGS.server, protocols=['iconograph-slave'], ssl_options=ssl_options)
+  client = Client(
+      FLAGS.config,
+      'wss://%s/ws/slave' % FLAGS.server,
+      protocols=['iconograph-slave'],
+      ssl_options=ssl_options)
   client.Loop()