Add support for timestamps in reboot commands

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'))

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)