From 388272c0f8b8e82ff1567ba7a58966bc203f9115 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <flamingcow@gmail.com>
Date: Wed, 11 May 2016 18:56:00 +0000
Subject: [PATCH] Add support for timestamps in reboot commands

---
 client/client.py  | 16 +++++++++++++---
 client/fetcher.py | 13 +++++++++++--
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/client/client.py b/client/client.py
index 6b01a3b..86d85c5 100755
--- a/client/client.py
+++ b/client/client.py
@@ -97,14 +97,17 @@ class Client(threadedclient.WebSocketClient):
       return
     self._OnNewManifest2()
 
-  def _OnNewManifest2(self):
-    fetch = fetcher.Fetcher(
+  def _GetFetcher(self):
+    return 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)
+
+  def _OnNewManifest2(self):
+    fetch = self._GetFetcher()
     fetch.Fetch()
     fetch.DeleteOldImages(skip={'%d.iso' % self._config['timestamp']})
 
@@ -115,7 +118,14 @@ class Client(threadedclient.WebSocketClient):
 
   def _OnCommand(self, data):
     if data['command'] == 'reboot':
-      subprocess.check_call(['reboot'])
+      self._OnReboot(data)
+
+  def _OnReboot(self, data):
+    if data['timestamp']:
+      fetch = self._GetFetcher()
+      fetch.Fetch(data['timestamp'])
+
+    subprocess.check_call(['reboot'])
 
   def received_message(self, msg):
     parsed = json.loads(msg.data.decode('utf8'))
diff --git a/client/fetcher.py b/client/fetcher.py
index 6b566c7..6c88ac5 100755
--- a/client/fetcher.py
+++ b/client/fetcher.py
@@ -107,6 +107,12 @@ class Fetcher(object):
     with open(path, 'w') as fh:
       json.dump(new_manifest, fh, indent=4)
 
+  def _FindImage(self, manifest, timestamp):
+    for image in manifest['images']:
+      if image['timestamp'] == timestamp:
+        return image
+    raise NoValidImage
+
   def _ChooseImage(self, manifest):
     hostname = socket.gethostname()
     hash_base = hashlib.sha256(hostname.encode('ascii'))
@@ -160,9 +166,12 @@ class Fetcher(object):
     os.symlink(filename, temp_path)
     os.rename(temp_path, current_path)
 
-  def Fetch(self):
+  def Fetch(self, force_timestamp=None):
     manifest = self._GetManifest()
-    image = self._ChooseImage(manifest)
+    if force_timestamp:
+      image = self._FindImage(manifest, timestamp)
+    else:
+      image = self._ChooseImage(manifest)
     self._FetchImage(image)
     self._SetCurrent(image)