Improve zip command

pupy/modules/zip.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 from pupylib.PupyModule import *
-from pupylib.utils.rpyc_utils import redirected_stdio
+from pupylib.utils.rpyc_utils import obtain
 
 __class_name__="Zip"
 
@@ -14,14 +14,37 @@ class Zip(PupyModule):
         self.arg_parser = PupyArgumentParser(prog="zip", description=self.__doc__)
         self.arg_parser.add_argument('source', type=str, help='path of the source file or directory to zip')
 
+        self.arg_parser.add_argument('-l', action='store_true', help='list file  (default: zip file)')
         self.arg_parser.add_argument('-u', action='store_true', help='unzip file (default: zip file)')
         self.arg_parser.add_argument('-d', dest='destination', help='path of the destination file (default: current directory)')
 
+    def nice_size(self, value):
+        if value > 1024*1024*1024:
+            return '{}G'.format(value/(1024*1024*1024))
+        elif value > 1024*1024:
+            return '{}M'.format(value/(1024*1024))
+        elif value > 1024:
+            return '{}K'.format(value/1024)
+        else:
+            return '{}B'.format(value)
+
     def run(self, args):
-        with redirected_stdio(self):
+        if args.l:
-            # zip
+            result, data = self.client.conn.modules["pupyutils.zip"].list(args.source)
-            if not args.u:
+            if result:
-                self.client.conn.modules["pupyutils.zip"].zip(args.source, args.destination)
+                data = obtain(data)
-            # unzip
+                log = args.source + ':\n' + '\n'.join(
+                    '{:>8} {}'.format(self.nice_size(file_size), filename) for filename, file_size in data
+                )
             else:
-                self.client.conn.modules["pupyutils.zip"].unzip(args.source, args.destination)
+                log = data
+
+        elif not args.u:
+            result, log = self.log(self.client.conn.modules["pupyutils.zip"].zip(args.source, args.destination))
+        else:
+            result, log = self.log(self.client.conn.modules["pupyutils.zip"].unzip(args.source, args.destination))
+
+        if result:
+            self.success(log)
+        else:
+            self.error(log)

pupy/packages/all/pupyutils/zip.py

@@ -17,8 +17,7 @@ def zip(src, dst):
     dst = try_unicode(dst)
 
     if not os.path.exists(src):
-        print "[-] The file \"%s\" does not exists" % src
+        return False, "The file \"%s\" does not exists" % src
-        return
 
     isDir = False
     if os.path.isdir(src):
@@ -33,8 +32,7 @@ def zip(src, dst):
 
     # To not overwrite an existing file
     if os.path.exists(dst):
-        print "[-] The destination file \"%s\" already exists" % dst
+        return False, "The destination file \"%s\" already exists" % dst
-        return
 
     # Zip process
     zf = zipfile.ZipFile(dst, "w", zipfile.ZIP_DEFLATED)
@@ -48,8 +46,8 @@ def zip(src, dst):
     else:
         zf.write(src)
 
-    print "[+] File zipped correctly: \"%s\"" % dst
     zf.close()
+    return True, "File zipped correctly: \"%s\"" % dst
 
 
 def unzip(src, dst):
@@ -57,8 +55,7 @@ def unzip(src, dst):
     dst = try_unicode(dst)
 
     if not os.path.exists(src):
-        print "[-] The file \"%s\" does not exists" % src
+        return False, "The file \"%s\" does not exists" % src
-        return
 
     if not dst:
         d = src.split(os.sep)
@@ -66,12 +63,27 @@ def unzip(src, dst):
 
     # To not overwrite an existing file
     if os.path.exists(dst):
-        print "[-] The destination file \"%s\" already exists" % dst
+        return False, "The destination file \"%s\" already exists" % dst
-        return
 
     if zipfile.is_zipfile(src):
         with zipfile.ZipFile(src, "r") as z:
             z.extractall(dst)
-        print "[+] File unzipped correctly: \"%s\"" % dst
+        return True, "File unzipped correctly: \"%s\"" % dst
     else:
-        print '[-] The zipped file does not have a valid zip format: \"%s\"'
+        return False, 'The zipped file does not have a valid zip format: \"%s\"' % src
+
+
+def list(src):
+    src = try_unicode(src)
+    if not os.path.exists(src):
+        return False, "The file \"%s\" does not exists" % src
+
+    if zipfile.is_zipfile(src):
+        with zipfile.ZipFile(src, "r") as z:
+            return True, [
+                (
+                    i.filename, i.file_size,
+                ) for i in z.infolist()
+            ]
+    else:
+        return False, 'The zipped file does not have a valid zip format: \"%s\"' % src