Add support for Quay's vulnerability tool

util/migrate/backfill_checksums.py

@@ -0,0 +1,67 @@
+import logging
+from app import storage as store
+from data.database import ImageStorage, ImageStoragePlacement, ImageStorageLocation, JOIN_LEFT_OUTER
+from digest import checksums
+
+logger = logging.getLogger(__name__)
+
+def _get_imagestorages_with_locations(query_modifier):
+  query = (ImageStoragePlacement
+           .select(ImageStoragePlacement, ImageStorage, ImageStorageLocation)
+           .join(ImageStorageLocation)
+           .switch(ImageStoragePlacement)
+           .join(ImageStorage, JOIN_LEFT_OUTER))
+  query = query_modifier(query)
+
+  location_list = list(query)
+
+  storages = {}
+  for location in location_list:
+    storage = location.storage
+
+    if not storage.id in storages:
+      storages[storage.id] = storage
+      storage.locations = set()
+    else:
+      storage = storages[storage.id]
+
+    storage.locations.add(location.location.name)
+
+  return storages.values()
+
+def backfill_checksum(imagestorage_with_locations):
+  try:
+    json_data = store.get_content(imagestorage_with_locations.locations, store.image_json_path(imagestorage_with_locations.uuid))
+    with store.stream_read_file(imagestorage_with_locations.locations, store.image_layer_path(imagestorage_with_locations.uuid)) as fp:
+      imagestorage_with_locations.checksum = 'sha256:{0}'.format(checksums.sha256_file(fp, json_data + '\n'))
+      imagestorage_with_locations.save()
+  except IOError as e:
+    if str(e).startswith("No such key"):
+      imagestorage_with_locations.checksum = 'unknown:{0}'.format(imagestorage_with_locations.uuid)
+      imagestorage_with_locations.save()
+  except:
+    logger.exception('exception when backfilling checksum of %s', imagestorage_with_locations.uuid)
+
+def backfill_checksums():
+  logger.setLevel(logging.DEBUG)
+
+  logger.debug('backfill_checksums: Starting')
+  logger.debug('backfill_checksums: This can be a LONG RUNNING OPERATION. Please wait!')
+
+  def limit_to_empty_checksum(query):
+    return query.where(ImageStorage.checksum >> None, ImageStorage.uploading == False).limit(100)
+
+  while True:
+    storages = _get_imagestorages_with_locations(limit_to_empty_checksum)
+    if len(storages) == 0:
+      logger.debug('backfill_checksums: Completed')
+      return
+
+    for storage in storages:
+      backfill_checksum(storage)
+
+if __name__ == "__main__":
+  logging.basicConfig(level=logging.DEBUG)
+  logging.getLogger('peewee').setLevel(logging.CRITICAL)
+  logging.getLogger('boto').setLevel(logging.CRITICAL)
+  backfill_checksums()

util/migrate/backfill_parent_id.py

@@ -1,38 +1,46 @@
 import logging
-
-from data.database import Image, ImageStorage, db, db_for_update
+from data.database import Image, ImageStorage, db
 from app import app
-from util.migrate import yield_random_entries
-
 
 logger = logging.getLogger(__name__)
 
-
 def backfill_parent_id():
   logger.setLevel(logging.DEBUG)
 
   logger.debug('backfill_parent_id: Starting')
   logger.debug('backfill_parent_id: This can be a LONG RUNNING OPERATION. Please wait!')
 
-  def fetch_batch():
-    return (Image
-            .select(Image.id, Image.ancestors)
-            .join(ImageStorage)
-            .where(Image.parent >> None, Image.ancestors != '/',
-                   ImageStorage.uploading == False))
+  # Check for any images without parent
+  has_images = bool(list(Image
+                         .select(Image.id)
+                         .join(ImageStorage)
+                         .where(Image.parent >> None, Image.ancestors != '/', ImageStorage.uploading == False)
+                         .limit(1)))
 
-  for to_backfill in yield_random_entries(fetch_batch, 10000, 0.3):
-    with app.config['DB_TRANSACTION_FACTORY'](db):
-      try:
-        image = db_for_update(Image
-                              .select()
-                              .where(Image.id == to_backfill.id)).get()
-        image.parent = to_backfill.ancestors.split('/')[-2]
-        image.save()
-      except Image.DoesNotExist:
-        pass
+  if not has_images:
+    logger.debug('backfill_parent_id: No migration needed')
+    return
 
-  logger.debug('backfill_parent_id: Completed')
+  while True:
+    # Load the record from the DB.
+    batch_images_ids = list(Image
+                            .select(Image.id)
+                            .join(ImageStorage)
+                            .where(Image.parent >> None, Image.ancestors != '/', ImageStorage.uploading == False)
+                            .limit(100))
+
+    if len(batch_images_ids) == 0:
+      logger.debug('backfill_parent_id: Completed')
+      return
+
+    for image_id in batch_images_ids:
+      with app.config['DB_TRANSACTION_FACTORY'](db):
+        try:
+          image = Image.select(Image.id, Image.ancestors).where(Image.id == image_id).get()
+          image.parent = image.ancestors.split('/')[-2]
+          image.save()
+        except Image.DoesNotExist:
+          pass
 
 if __name__ == "__main__":
   logging.basicConfig(level=logging.DEBUG)