updated matrix program

matrix.py

@@ -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+")
+    i = 0
+    with open(backup_file, "w+") as file_desc:
         result = []
-    for rd,dr,fs in os.walk(directory):
+        for root_dir, dirs, files in os.walk(directory):
-        for d in dr:
+            for d in dirs:
-            res = os.path.join(rd, d)
+                res = os.path.join(root_dir, d)
                 result.append(res)
-        for f in fs:
+            for f in files:
-            res = os.path.join(rd, f)
+                res = os.path.join(root_dir, f)
                 result.append(res)
-    for f_name_tmp in result:
+        for f_name in result:
             try:
-            f_name = f_name_tmp.replace(directory,'').lstrip("/")
+                i += 1
-            f_perm = str(oct(os.stat(f_name_tmp).st_mode))
+                f_uid = os.stat(f_name).st_uid
-            line = "%s|%s\n" % (f_name,f_perm)
+                f_gid = os.stat(f_name).st_gid
-            fd.write(line)
+                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
-    fd.close() 
+        print("Backup {} files permissions, storing to {}".format(i, backup_file))
 
 def restore():
-    fd = open(backup_file,"r+")
+    i = 0
-    for f in fd:
+    with open(backup_file, "r+") as file_desc:
+        for f in file_desc:
             line = f.strip("\n")
-        f_name_tmp = line.split('|')[0]
+            f_name = line.split('|')[0]
-        f_name = "%s/%s" % (directory,f_name_tmp)
+            f_uid = line.split('|')[1]
-        f_perm = 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()
-