feat (cvedb): rollback cachedir if cvedb refresh fails (#1225)

* feat (cvedb): rollback cachedir if cvedb refresh fails

* pass backup_cachedir argument in cvedb

* update test_cli.py log assertion

Author: imsahil007

cve_bin_tool/cli.py

@@ -308,6 +308,8 @@ def main(argv=None):
         with ErrorHandler(mode=error_mode, logger=LOGGER):
             raise CVEDataMissing("No data in CVE Database")
 
+    cvedb_orig.remove_cache_backup()
+
     # Input validation
     if not args["directory"] and not args["input_file"] and not args["package_list"]:
         parser.print_usage()

cve_bin_tool/cvedb.py

@@ -36,6 +36,9 @@
 logging.basicConfig(level=logging.DEBUG)
 
 # database defaults
+DISK_LOCATION_BACKUP = os.path.join(
+    os.path.expanduser("~"), ".cache", "cve-bin-tool-backup"
+)
 DISK_LOCATION_DEFAULT = os.path.join(os.path.expanduser("~"), ".cache", "cve-bin-tool")
 DBNAME = "cve.db"
 OLD_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "cvedb")
@@ -47,6 +50,7 @@ class CVEDB:
     """
 
     CACHEDIR = DISK_LOCATION_DEFAULT
+    BACKUPCACHEDIR = DISK_LOCATION_BACKUP
     FEED = "https://nvd.nist.gov/vuln/data-feeds"
     LOGGER = LOGGER.getChild("CVEDB")
     NVDCVE_FILENAME_TEMPLATE = "nvdcve-1.1-{}.json.gz"
@@ -59,12 +63,16 @@ def __init__(
         self,
         feed=None,
         cachedir=None,
+        backup_cachedir=None,
         version_check=True,
         session=None,
         error_mode=ErrorMode.TruncTrace,
     ):
         self.feed = feed if feed is not None else self.FEED
         self.cachedir = cachedir if cachedir is not None else self.CACHEDIR
+        self.backup_cachedir = (
+            backup_cachedir if backup_cachedir is not None else self.BACKUPCACHEDIR
+        )
         self.error_mode = error_mode
         # Will be true if refresh was successful
         self.was_updated = False
@@ -78,6 +86,9 @@ def __init__(
         self.session = session
         self.cve_count = -1
 
+        if not os.path.exists(self.dbpath):
+            self.rollback_cache_backup()
+
     def get_cve_count(self):
         if self.cve_count == -1:
             # Force update
@@ -618,8 +629,9 @@ def curl_versions(self):
         ]
 
     def clear_cached_data(self):
+        self.create_cache_backup()
         if os.path.exists(self.cachedir):
-            self.LOGGER.warning(f"Deleting cachedir {self.cachedir}")
+            self.LOGGER.warning(f"Updating cachedir {self.cachedir}")
             shutil.rmtree(self.cachedir)
         # Remove files associated with pre-1.0 development tree
         if os.path.exists(OLD_CACHE_DIR):
@@ -636,3 +648,29 @@ def db_close(self):
         if self.connection:
             self.connection.close()
             self.connection = None
+
+    def create_cache_backup(self):
+        """Creates a backup of the cachedir in case anything fails"""
+        if os.path.exists(self.cachedir):
+            self.LOGGER.debug(
+                f"Creating backup of cachedir {self.cachedir} at {self.backup_cachedir}"
+            )
+            self.remove_cache_backup()
+            shutil.copytree(self.cachedir, self.backup_cachedir)
+
+    def remove_cache_backup(self):
+        """Removes the backup if database was successfully loaded"""
+        if os.path.exists(self.backup_cachedir):
+            self.LOGGER.debug(f"Removing backup cache from {self.backup_cachedir}")
+            shutil.rmtree(self.backup_cachedir)
+
+    def rollback_cache_backup(self):
+        """Rollback the cachedir backup in case anything fails"""
+        if os.path.exists(os.path.join(self.backup_cachedir, DBNAME)):
+            self.LOGGER.info(f"Rolling back the cache to its previous state")
+            if os.path.exists(self.cachedir):
+                shutil.rmtree(self.cachedir)
+            shutil.move(self.backup_cachedir, self.cachedir)
+
+    def __del__(self):
+        self.rollback_cache_backup()

test/test_cli.py

@@ -206,7 +206,7 @@ def test_update(self, caplog):
         assert (
             "cve_bin_tool.CVEDB",
             logging.WARNING,
-            f"Deleting cachedir {db_path}",
+            f"Updating cachedir {db_path}",
         ) in caplog.record_tuples and (
             "cve_bin_tool.CVEDB",
             logging.INFO,