More reliable cleanup

client/fetcher.py

@@ -130,8 +130,8 @@ class Fetcher(object):
     resp = self._session.get(url, stream=True)
 
     hash_obj = hashlib.sha256()
+    with tempfile.NamedTemporaryFile(dir=self._image_dir, delete=False) as fh:
       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)

client/update_grub.py

@@ -32,8 +32,8 @@ 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:
+      try:
         current = os.readlink(os.path.join(self._image_dir, 'current'))
 
         fh.write("""
@@ -68,3 +68,6 @@ menuentry "%(image_filename)s (%(volume_id)s)" --hotkey=%(hotkey)s {
 
         fh.flush()
         os.rename(fh.name, os.path.join(grub_dir, 'grub.cfg'))
+      except:
+        os.unlink(fh.name)
+        raise