diff --git a/client/client.py b/client/client.py index 49b97d5..6b01a3b 100755 --- a/client/client.py +++ b/client/client.py @@ -86,7 +86,7 @@ class Client(threadedclient.WebSocketClient): return int(float(fh.readline().split(' ', 1)[0])) def _NextTimestamp(self): - next_image = os.path.basename(os.readlink('/isodevice/iconograph/current')) + next_image = lib.GetCurrentImage() return int(next_image.split('.', 1)[0]) def _OnImageTypes(self, data): @@ -106,7 +106,7 @@ class Client(threadedclient.WebSocketClient): FLAGS.https_client_cert, FLAGS.https_client_key) fetch.Fetch() - fetch.DeleteOldImages() + fetch.DeleteOldImages(skip={'%d.iso' % self._config['timestamp']}) update = update_grub.GrubUpdater( FLAGS.image_dir, diff --git a/client/fetcher.py b/client/fetcher.py index f8f6b28..6b566c7 100755 --- a/client/fetcher.py +++ b/client/fetcher.py @@ -166,9 +166,10 @@ class Fetcher(object): self._FetchImage(image) self._SetCurrent(image) - def DeleteOldImages(self, max_images=5): + def DeleteOldImages(self, max_images=5, skip=None): if not max_images: return + skip = skip or set() images = [] for filename in os.listdir(self._image_dir): match = self._FILE_REGEX.match(filename) @@ -177,6 +178,8 @@ class Fetcher(object): images.append((int(match.group('timestamp')), filename)) images.sort(reverse=True) for timestamp, filename in images[max_images:]: + if filename in skip: + continue print('Deleting old image:', filename, flush=True) path = os.path.join(self._image_dir, filename) os.unlink(path) diff --git a/client/lib.py b/client/lib.py index abffe26..50e9ee5 100644 --- a/client/lib.py +++ b/client/lib.py @@ -1,3 +1,4 @@ +import os import re import subprocess @@ -13,3 +14,8 @@ def GetVolumeID(path): ]) match = _VOLUME_ID_REGEX.search(isoinfo) return match.group('volume_id').decode('ascii') + + +def GetCurrentImage(image_dir='/isodevice/iconograph'): + current_path = os.path.join(image_dir, 'current') + return os.path.basename(os.readlink(current_path)) diff --git a/client/update_grub.py b/client/update_grub.py index 7afcc2d..c10fc96 100755 --- a/client/update_grub.py +++ b/client/update_grub.py @@ -23,13 +23,13 @@ class GrubUpdater(object): with tempfile.NamedTemporaryFile('w', dir=grub_dir, delete=False) as fh: try: - current = os.readlink(os.path.join(self._image_dir, 'current')) + current = lib.GetCurrentImage(self._image_dir) fh.write(""" set timeout=5 set default=%(default_image_filename)s """ % { - 'default_image_filename': os.path.basename(current), + 'default_image_filename': current, }) files = []