Added support to backup folders during upgrade.

The upgrade framework has been updated to support backup and restore
operations for folders and their contents.

Ticket #583

2 files changed, 137 insertions, 26 deletions

diff --git a/base/common/python/pki/upgrade.py b/base/common/python/pki/upgrade.py
index 675f3b903..d1cda6bac 100644
--- a/base/common/python/pki/upgrade.py
+++ b/base/common/python/pki/upgrade.py
@@ -278,7 +278,7 @@ class PKIUpgradeScriptlet(object):
         # Callback method to upgrade the system.
         pass
 
-    def upgrade(self):
+    def init(self):
 
         backup_dir = self.get_backup_dir()
 
@@ -289,6 +289,8 @@ class PKIUpgradeScriptlet(object):
         # create backup dir
         os.makedirs(backup_dir)
 
+    def upgrade(self):
+
         try:
             if not self.can_upgrade():
                 if verbose: print 'Skipping system.'
@@ -324,46 +326,93 @@ class PKIUpgradeScriptlet(object):
         if os.path.exists(oldfiles):
 
             # restore all backed up files
-            for root, _, filenames in os.walk(oldfiles):  #unused item _ for dirnames
-                path = root[len(oldfiles):]
+            for sourcepath, _, filenames in os.walk(oldfiles):  #unused item _ for dirnames
+
+                destpath = sourcepath[len(oldfiles):]
+                if destpath == '': destpath = '/'
+
+                if not os.path.isdir(destpath):
+                    if verbose: print 'Restoring ' + destpath
+                    pki.util.copydirs(sourcepath, destpath)
+
                 for filename in filenames:
-                    source = root + '/' + filename
-                    target = path + '/' + filename
+                    sourcefile = os.path.join(sourcepath, filename)
+                    targetfile = os.path.join(destpath, filename)
 
-                    if verbose: print 'Restoring ' + target
-                    pki.util.copyfile(source, target)
+                    if verbose: print 'Restoring ' + targetfile
+                    pki.util.copyfile(sourcefile, targetfile)
 
         newfiles = backup_dir + '/newfiles'
         if os.path.exists(newfiles):
 
-            # remove files that did not exist before upgrade
+            # get paths that did not exist before upgrade
+            paths = []
             with open(newfiles, 'r') as f:
-                for filename in f:
-                    filename = filename.strip('\n')
+                for path in f:
+                    path = path.strip('\n')
+                    paths.append(path)
 
-                    if os.path.exists(filename):
-                        if verbose: print 'Deleting ' + filename
-                        os.remove(filename)
+            # remove paths in reverse order
+            paths.reverse()
+            for path in paths:
 
-    def backup(self, filename):
+                if not os.path.exists(path): continue
+                if verbose: print 'Deleting ' + path
+
+                if os.path.isfile(path):
+                    os.remove(path)
+                else:
+                    shutil.rmtree(path)
+
+    def backup(self, path):
 
         backup_dir = self.get_backup_dir()
-        backup_file = backup_dir + '/oldfiles' + filename
 
-        if os.path.exists(filename):
+        if not os.path.exists(backup_dir):
+            os.makedirs(backup_dir)
+
+        if os.path.exists(path):
+
+            # if path exists, keep a copy
+
+            oldfiles = backup_dir + '/oldfiles'
+            if not os.path.exists(oldfiles):
+                os.mkdir(oldfiles)
+
+            dest = oldfiles + path
+
+            sourceparent = os.path.dirname(path)
+            destparent = os.path.dirname(dest)
+
+            pki.util.copydirs(sourceparent, destparent)
+
+            if os.path.isfile(path):
+                if verbose: print 'Saving ' + path
+                pki.util.copyfile(path, dest)
+
+            else:
+                for sourcepath, _, filenames in os.walk(path):
+
+                    relpath = sourcepath[len(path):]
+                    destpath = dest + relpath
+
+                    if verbose: print 'Saving ' + sourcepath
+                    pki.util.copydirs(sourcepath, destpath)
 
-            # if file exists, keep a copy
+                    for filename in filenames:
+                        sourcefile = os.path.join(sourcepath, filename)
+                        targetfile = os.path.join(destpath, filename)
 
-            if verbose: print 'Saving ' + filename
-            pki.util.copyfile(filename, backup_file)
+                        if verbose: print 'Saving ' + sourcefile
+                        pki.util.copyfile(sourcefile, targetfile)
 
         else:
 
-            # otherwise, keep the name
+            # otherwise, record the name
 
-            if verbose: print 'Recording ' + filename
+            if verbose: print 'Recording ' + path
             with open(backup_dir + '/newfiles', 'a') as f:
-                f.write(filename + '\n')
+                f.write(path + '\n')
 
     def __eq__(self, other):
         return self.version == other.version and self.index == other.index
@@ -543,6 +592,7 @@ class PKIUpgrader(object):
                     raise pki.PKIException('Upgrade canceled.')
 
             try:
+                scriptlet.init()
                 scriptlet.upgrade()
 
             except pki.PKIException as e:
diff --git a/base/common/python/pki/util.py b/base/common/python/pki/util.py
index 9b501a342..4d25390a6 100644
--- a/base/common/python/pki/util.py
+++ b/base/common/python/pki/util.py
@@ -22,13 +22,74 @@
 import os
 import shutil
 
+def copy(source, dest):
+    """
+    Copy a file or a folder and its contents.
+    """
+
+    # remove trailing slashes
+    if source[-1] == '/': source = source[:-1]
+    if dest[-1] == '/': dest = dest[:-1]
+
+    sourceparent = os.path.dirname(source)
+    destparent = os.path.dirname(dest)
+
+    copydirs(sourceparent, destparent)
+
+    if os.path.isfile(source):
+        copyfile(source, dest)
+
+    else:
+        for sourcepath, _, filenames in os.walk(source):
+
+            relpath = sourcepath[len(source):]
+            destpath = dest + relpath
+            if destpath == '': destpath = '/'
+
+            copydirs(sourcepath, destpath)
+
+            for filename in filenames:
+                sourcefile = os.path.join(sourcepath, filename)
+                targetfile = os.path.join(destpath, filename)
+                copyfile(sourcefile, targetfile)
 
 def copyfile(source, dest):
+    """
+    Copy a file or link while preserving its attributes.
+    """
+
+    if os.path.islink(source):
+        target = os.readlink(source)
+        os.symlink(target, dest)
+
+        st = os.lstat(source)
+        os.lchown(dest, st.st_uid, st.st_gid)
+
+    else:
+        shutil.copyfile(source, dest)
+
+        st = os.stat(source)
+        os.utime(dest, (st.st_atime, st.st_mtime))
+        os.chmod(dest, st.st_mode)
+        os.chown(dest, st.st_uid, st.st_gid)
+
+def copydirs(source, dest):
+    """
+    Copy a folder and its parents while preserving their attributes.
+    """
+
+    if os.path.exists(dest):
+        return
+
+    destparent = os.path.dirname(dest)
+
+    if not os.path.exists(destparent):
+        sourceparent = os.path.dirname(source)
+        copydirs(sourceparent, destparent)
 
-    dest_dir = os.path.dirname(dest)
-    if not os.path.exists(dest_dir):
-        os.makedirs(dest_dir)
+    os.mkdir(dest)
 
-    shutil.copy2(source, dest)
     st = os.stat(source)
+    os.utime(dest, (st.st_atime, st.st_mtime))
+    os.chmod(dest, st.st_mode)
     os.chown(dest, st.st_uid, st.st_gid)