Reworked the backup/restore to properly handle the cache directory

inclusion.

SickBeard.py

@@ -264,10 +264,10 @@ class SickRage(object):
         # Check if we need to perform a restore first
         restoreDir = os.path.join(sickbeard.DATA_DIR, 'restore')
         if self.consoleLogging and os.path.exists(restoreDir):
-            if self.restore(restoreDir, sickbeard.DATA_DIR):
-                sys.stdout.write("Restore successful...\n")
+            if self.restoreDB(restoreDir, sickbeard.DATA_DIR):
+                sys.stdout.write("Restore: restoring DB and config.ini successful...\n")
             else:
-                sys.stdout.write("Restore FAILED!\n")
+                sys.stdout.write("Restore: restoring DB and config.ini FAILED!\n")
 
         # Load the config and publish it to the sickbeard package
         if self.consoleLogging and not os.path.isfile(sickbeard.CONFIG_FILE):
@@ -447,16 +447,17 @@ class SickRage(object):
                     logger.ERROR)
                 logger.log(traceback.format_exc(), logger.DEBUG)
 
-    def restore(self, srcDir, dstDir):
+    def restoreDB(self, srcDir, dstDir):
         try:
-            for file in os.listdir(srcDir):
-                srcFile = os.path.join(srcDir, file)
-                dstFile = os.path.join(dstDir, file)
-                bakFile = os.path.join(dstDir, file + '.bak')
+            filesList = ['sickbeard.db', 'config.ini']
+
+            for filename in filesList:
+                srcFile = os.path.join(srcDir, filename)
+                dstFile = os.path.join(dstDir, filename)
+                bakFile = os.path.join(dstDir, '{0}.bak-{1}'.format(filename, datetime.datetime.strftime(datetime.datetime.now(), '%Y%m%d_%H%M%S')))
+                if os.path.isfile(dstFile):
                     shutil.move(dstFile, bakFile)
                 shutil.move(srcFile, dstFile)
-
-            os.rmdir(srcDir)
             return True
         except:
             return False

sickbeard/__init__.py

@@ -631,6 +631,31 @@ def initialize(consoleLogging=True):
             logger.log(u"!!! Creating local cache dir failed, using system default", logger.ERROR)
             CACHE_DIR = None
 
+        # Check if we need to perform a restore of the cache folder
+        restoreDir = os.path.join(DATA_DIR, 'restore')
+        if os.path.exists(restoreDir):
+            def restoreCache(srcDir, dstDir):
+                import ntpath
+                import shutil
+
+                def path_leaf(path):
+                    head, tail = ntpath.split(path)
+                    return tail or ntpath.basename(head)
+
+                try:
+                    if os.path.isdir(dstDir):
+                        bakFilename = '{0}-{1}'.format(path_leaf(dstDir), datetime.datetime.strftime(datetime.datetime.now(), '%Y%m%d_%H%M%S'))
+                        shutil.move(dstDir, os.path.join(ntpath.dirname(dstDir), bakFilename))
+
+                    shutil.move(srcDir, dstDir)
+                    logger.log(u"Restore: restoring cache successful", logger.INFO)
+                except Exception as e:
+                    logger.log(u"Restore: restoring cache failed", logger.ERROR)
+                finally:
+                    os.rmdir(restoreDir)
+
+            restoreCache(os.path.join(restoreDir, 'cache'), CACHE_DIR)
+
         # clean cache folders
         if CACHE_DIR:
             helpers.clearCache()

sickbeard/helpers.py

@@ -1067,6 +1067,34 @@ def extractZip(archive, targetDir):
         return False
 
 
+def backupConfigZip(fileList, archive, arcname = None):
+    try:
+        a = zipfile.ZipFile(archive, 'w', zipfile.ZIP_DEFLATED)
+        for f in fileList:
+            a.write(f, os.path.relpath(f, arcname))
+        a.close()
+        return True
+    except Exception as e:
+        logger.log(u"Zip creation error: " + str(e), logger.ERROR)
+        return False
+
+
+def restoreConfigZip(archive, targetDir):
+    try:
+        if not os.path.exists(targetDir):
+            os.mkdir(targetDir)
+
+        zip_file = zipfile.ZipFile(archive, 'r')
+        for member in zip_file.namelist():
+            zip_file.extract(member, targetDir)
+        zip_file.close()
+        return True
+    except Exception as e:
+        logger.log(u"Zip extraction error: " + str(e), logger.ERROR)
+        shutil.rmtree(targetDir)
+        return False
+
+
 def mapIndexersToShow(showObj):
     mapped = {}
 

sickbeard/webserve.py

@@ -3599,11 +3599,14 @@ class ConfigBackupRestore(Config):
             source = [os.path.join(sickbeard.DATA_DIR, 'sickbeard.db'), sickbeard.CONFIG_FILE]
             target = os.path.join(backupDir, 'sickrage-' + time.strftime('%Y%m%d%H%M%S') + '.zip')
 
-            for (dir, _, files) in os.walk(sickbeard.CACHE_DIR):
-                for f in files:
-                    source.append(os.path.join(dir, f))
+            for (path, dirs, files) in os.walk(sickbeard.CACHE_DIR, topdown=True):
+                for dirname in dirs:
+                    if path == sickbeard.CACHE_DIR and dirname not in ['images']:
+                        dirs.remove(dirname)
+                for filename in files:
+                    source.append(os.path.join(path, filename))
 
-            if helpers.makeZip(source, target):
+            if helpers.backupConfigZip(source, target, sickbeard.DATA_DIR):
                 finalResult += "Successful backup to " + target
             else:
                 finalResult += "Backup FAILED"
@@ -3623,7 +3626,7 @@ class ConfigBackupRestore(Config):
             source = backupFile
             target_dir = os.path.join(sickbeard.DATA_DIR, 'restore')
 
-            if helpers.extractZip(source, target_dir):
+            if helpers.restoreConfigZip(source, target_dir):
                 finalResult += "Successfully extracted restore files to " + target_dir
                 finalResult += "<br>Restart sickrage to complete the restore."
             else: