updated matrix program

matrix.py

@@ -15,7 +15,7 @@ backup_file = ""
 rep = ""
 
 parser = argparse.ArgumentParser(description='Parse arguments')
-parser.add_argument('mode', metavar='mode', type=str, choices=['backup','restore'], help='Mode')
+parser.add_argument('mode', metavar='mode', type=str, choices=['backup', 'restore'], help='Mode')
 parser.add_argument('backup_file', metavar='backup_file', type=str, help='Backup file')
 parser.add_argument('directory', metavar='directory', type=str, help='Directory')
 
@@ -23,42 +23,50 @@ args = parser.parse_args()
 
 mode = args.mode
 backup_file = args.backup_file
-directory = args.directory.rstrip("/")
+directory = args.directory
 
 def backup():
-    fd = open(backup_file,"w+")
-    result=[]
-    for rd,dr,fs in os.walk(directory):
-        for d in dr:
-            res = os.path.join(rd, d)
-            result.append(res)
-        for f in fs:
-            res = os.path.join(rd, f)
-            result.append(res)
-    for f_name_tmp in result:
-        try:
-            f_name = f_name_tmp.replace(directory,'').lstrip("/")
-            f_perm = str(oct(os.stat(f_name_tmp).st_mode))
-            line = "%s|%s\n" % (f_name,f_perm)
-            fd.write(line)
-        except OSError:
-            pass
-    fd.close() 
+    i = 0
+    with open(backup_file, "w+") as file_desc:
+        result = []
+        for root_dir, dirs, files in os.walk(directory):
+            for d in dirs:
+                res = os.path.join(root_dir, d)
+                result.append(res)
+            for f in files:
+                res = os.path.join(root_dir, f)
+                result.append(res)
+        for f_name in result:
+            try:
+                i += 1
+                f_uid = os.stat(f_name).st_uid
+                f_gid = os.stat(f_name).st_gid
+                f_perm = str(oct(os.stat(f_name).st_mode))
+                line = "{}|{}|{}|{}\n".format(f_name, f_uid, f_gid, f_perm)
+                file_desc.write(line)
+            except OSError:
+                pass
+        print("Backup {} files permissions, storing to {}".format(i, backup_file))
 
 def restore():
-    fd = open(backup_file,"r+")
-    for f in fd:
-        line = f.strip("\n")
-        f_name_tmp = line.split('|')[0]
-        f_name = "%s/%s" % (directory,f_name_tmp)
-        f_perm = line.split('|')[1]
-        try:
-            os.chmod(f_name,int(f_perm,8))
-        except OSError:
-            pass
+    i = 0
+    with open(backup_file, "r+") as file_desc:
+        for f in file_desc:
+            line = f.strip("\n")
+            f_name = line.split('|')[0]
+            f_uid = line.split('|')[1]
+            f_gid = line.split('|')[2]
+            f_perm = line.split('|')[3]
+            try:
+                os.chmod(f_name, int(f_perm, 8))
+                os.chown(f_name, int(f_uid), int(f_gid))
+            except OSError:
+                pass
+            i += 1
+
+    print("Restored {} files permissions".format(i))
 
 if mode == "backup":
     backup()
 elif mode == "restore":
     restore()
-