From f43e9913e94413c5efa8ac8fc2ec180e9bb58f92 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <flamingcow@gmail.com>
Date: Tue, 10 May 2016 20:55:50 +0000
Subject: [PATCH] More reliable cleanup

---
 client/fetcher.py     | 22 +++++++++++-----------
 client/update_grub.py | 41 ++++++++++++++++++++++-------------------
 2 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/client/fetcher.py b/client/fetcher.py
index 02033b6..f8f6b28 100755
--- a/client/fetcher.py
+++ b/client/fetcher.py
@@ -130,17 +130,17 @@ class Fetcher(object):
     resp = self._session.get(url, stream=True)
 
     hash_obj = hashlib.sha256()
-    try:
-      fh = tempfile.NamedTemporaryFile(dir=self._image_dir, delete=False)
-      for data in resp.iter_content(self._BUF_SIZE):
-        hash_obj.update(data)
-        fh.write(data)
-      if hash_obj.hexdigest() != image['hash']:
-        raise InvalidHash
-      os.rename(fh.name, path)
-    except:
-      os.unlink(fh.name)
-      raise
+    with tempfile.NamedTemporaryFile(dir=self._image_dir, delete=False) as fh:
+      try:
+        for data in resp.iter_content(self._BUF_SIZE):
+          hash_obj.update(data)
+          fh.write(data)
+        if hash_obj.hexdigest() != image['hash']:
+          raise InvalidHash
+        os.rename(fh.name, path)
+      except:
+        os.unlink(fh.name)
+        raise
 
   def _SetCurrent(self, image):
     filename = '%d.iso' % (image['timestamp'])
diff --git a/client/update_grub.py b/client/update_grub.py
index a63f205..94cfbf3 100755
--- a/client/update_grub.py
+++ b/client/update_grub.py
@@ -32,25 +32,25 @@ class GrubUpdater(object):
   def Update(self):
     grub_dir = os.path.join(self._boot_dir, 'grub')
 
-    # TODO: clean up if we fail here
     with tempfile.NamedTemporaryFile('w', dir=grub_dir, delete=False) as fh:
-      current = os.readlink(os.path.join(self._image_dir, 'current'))
+      try:
+        current = os.readlink(os.path.join(self._image_dir, 'current'))
 
-      fh.write("""
+        fh.write("""
 set timeout=5
 set default=%(default_image_filename)s
 """ % {
-        'default_image_filename': os.path.basename(current),
-      })
+          'default_image_filename': os.path.basename(current),
+        })
 
-      files = []
-      for filename in os.listdir(self._image_dir):
-        if not filename.endswith('.iso'):
-          continue
-        files.append(filename)
+        files = []
+        for filename in os.listdir(self._image_dir):
+          if not filename.endswith('.iso'):
+            continue
+          files.append(filename)
 
-      for i, filename in enumerate(sorted(files, reverse=True)):
-        fh.write("""
+        for i, filename in enumerate(sorted(files, reverse=True)):
+          fh.write("""
 menuentry "%(image_filename)s (%(volume_id)s)" --hotkey=%(hotkey)s {
   search --no-floppy --file --set=root %(image_path)s/%(image_filename)s
   iso_path="%(image_path)s/%(image_filename)s"
@@ -60,11 +60,14 @@ menuentry "%(image_filename)s (%(volume_id)s)" --hotkey=%(hotkey)s {
   configfile /boot/grub/loopback.cfg
 }
 """ % {
-          'image_filename': filename,
-          'image_path': self._image_path,
-          'hotkey': self._HOTKEYS[i],
-          'volume_id': self._GetVolumeID(os.path.join(self._image_dir, filename)),
-        })
+            'image_filename': filename,
+            'image_path': self._image_path,
+            'hotkey': self._HOTKEYS[i],
+            'volume_id': self._GetVolumeID(os.path.join(self._image_dir, filename)),
+          })
 
-      fh.flush()
-      os.rename(fh.name, os.path.join(grub_dir, 'grub.cfg'))
+        fh.flush()
+        os.rename(fh.name, os.path.join(grub_dir, 'grub.cfg'))
+      except:
+        os.unlink(fh.name)
+        raise