Extractor rewritten

1 files changed, 115 insertions, 53 deletions

diff --git a/module/plugins/hooks/ExtractArchive.py b/module/plugins/hooks/ExtractArchive.py
index ddec8319b..16942bef0 100644
--- a/module/plugins/hooks/ExtractArchive.py
+++ b/module/plugins/hooks/ExtractArchive.py
@@ -51,32 +51,33 @@ if os.name != "nt":
     from pwd import getpwnam
 
 from module.plugins.Hook import Hook, threaded, Expose
-from module.plugins.internal.AbstractExtractor import ArchiveError, CRCError, WrongPassword
-from module.utils import save_join, fs_encode
+from module.plugins.internal.AbstractExtractor import ArchiveError, CRCError, PasswordError
+from module.utils import save_join, uniqify
 
 
 class ExtractArchive(Hook):
     __name__    = "ExtractArchive"
     __type__    = "hook"
-    __version__ = "0.20"
-
-    __config__ = [("activated", "bool", "Activated", True),
-                  ("fullpath", "bool", "Extract full path", True),
-                  ("overwrite", "bool", "Overwrite files", True),
-                  ("passwordfile", "file", "password file", "archive_password.txt"),
-                  ("deletearchive", "bool", "Delete archives when done", False),
-                  ("subfolder", "bool", "Create subfolder for each package", False),
-                  ("destination", "folder", "Extract files to", ""),
-                  ("excludefiles", "str", "Exclude files from unpacking (seperated by ;)", ""),
-                  ("recursive", "bool", "Extract archives in archvies", True),
-                  ("queue", "bool", "Wait for all downloads to be finished", True),
-                  ("renice", "int", "CPU Priority", 0)]
+    __version__ = "1.00"
+
+    __config__ = [("activated"    , "bool"  , "Activated"                                 , True                                                                     ),
+                  ("fullpath"     , "bool"  , "Extract full path"                         , True                                                                     ),
+                  ("overwrite"    , "bool"  , "Overwrite files"                           , False                                                                    ),
+                  ("keepbroken"   , "bool"  , "Extract broken archives"                   , False                                                                    ),
+                  ("repair"       , "bool"  , "Repair broken archives"                    , True                                                                     ),
+                  ("passwordfile" , "file"  , "Store passwords in file"                   , "archive_password.txt"                                                   ),
+                  ("delete"       , "bool"  , "Delete archive when successfully extracted", False                                                                    ),
+                  ("subfolder"    , "bool"  , "Create subfolder for each package"         , False                                                                    ),
+                  ("destination"  , "folder", "Extract files to"                          , ""                                                                       ),
+                  ("extensions"   , "str"   , "Extract the following extensions"          , "7z,bz2,bzip2,gz,gzip,lha,lzh,lzma,rar,tar,taz,tbz,tbz2,tgz,xar,xz,z,zip"),
+                  ("excludefiles" , "str"   , "Don't extract the following files"         , "*.nfo,*.DS_Store,index.dat,thumb.db"                                    ),
+                  ("recursive"    , "bool"  , "Extract archives in archives"              , True                                                                     ),
+                  ("queue"        , "bool"  , "Wait for all downloads to be finished"     , True                                                                     ),
+                  ("renice"       , "int"   , "CPU Priority"                              , 0                                                                        )]
 
     __description__ = """Extract different kind of archives"""
     __license__     = "GPLv3"
-    __authors__     = [("RaNaN", "ranan@pyload.org"),
-                       ("AndroKev", None),
-                       ("Walter Purcaro", "vuolter@gmail.com")]
+    __authors__     = [("Walter Purcaro", "vuolter@gmail.com")]
 
 
     event_list = ["allDownloadsProcessed"]
@@ -92,7 +93,7 @@ class ExtractArchive(Hook):
         self.passwords = []
         names = []
 
-        for p in ("UnRar", "UnZip"):
+        for p in ("UnRar", "SevenZip", "UnZip"):
             try:
                 module = self.core.pluginManager.loadModule("internal", p)
                 klass = getattr(module, p)
@@ -154,13 +155,21 @@ class ExtractArchive(Hook):
         extracted = []
         failed    = []
 
+        clearlist = lambda string: [x.lstrip('.') for x in string.replace(' ', '').replace(',', '|').replace(';', '|').split('|')]
+
         destination  = self.getConfig("destination")
         subfolder    = self.getConfig("subfolder")
         fullpath     = self.getConfig("fullpath")
         overwrite    = self.getConfig("overwrite")
-        excludefiles = self.getConfig("excludefiles")
+        extensions   = clearlist(self.getConfig("extensions"))
+        excludefiles = clearlist(self.getConfig("excludefiles"))
         renice       = self.getConfig("renice")
         recursive    = self.getConfig("recursive")
+        delete       = self.getConfig("delete")
+        keepbroken   = self.getConfig("keepbroken")
+
+        if extensions:
+            self.logDebug("Extensions allowed: %s" % "|.".join(extensions))
 
         # reload from txt file
         self.reloadPasswords()
@@ -171,7 +180,7 @@ class ExtractArchive(Hook):
         #iterate packages -> plugins -> targets
         for pid in ids:
             p = self.core.files.getPackage(pid)
-            self.logInfo(_("Check package %s") % p.name)
+            self.logInfo(_("Check package: %s") % p.name)
             if not p:
                 continue
 
@@ -179,21 +188,25 @@ class ExtractArchive(Hook):
             out = save_join(dl, p.folder, destination, "")  #: force trailing slash
 
             if subfolder:
-                out = save_join(out, fs_encode(p.folder))
+                out = save_join(out, p.folder)
 
             if not exists(out):
                 makedirs(out)
 
             files_ids = [(save_join(dl, p.folder, x['name']), x['id']) for x in p.getChildren().itervalues()]
-            matched = False
-            success = True
+            matched   = False
+            success   = True
 
             # check as long there are unseen files
             while files_ids:
                 new_files_ids = []
 
+                if extensions:
+                    files_ids = [(file, id) for file, id in files_ids if filter(lambda ext: file.endswith(ext), extensions)]
+
                 for plugin in self.plugins:
                     targets = plugin.getTargets(files_ids)
+
                     if targets:
                         self.logDebug("Targets for %s: %s" % (plugin.__name__, targets))
                         matched = True
@@ -205,19 +218,31 @@ class ExtractArchive(Hook):
 
                         processed.append(target)  # prevent extracting same file twice
 
-                        self.logInfo(basename(target), _("Extract to %s") % out)
+                        self.logInfo(basename(target), _("Extract to: %s") % out)
                         try:
-                            klass = plugin(self, target, out, fullpath, overwrite, excludefiles, renice)
+                            klass = plugin(self,
+                                           target,
+                                           out,
+                                           p.password,
+                                           fullpath,
+                                           overwrite,
+                                           excludefiles,
+                                           renice,
+                                           delete,
+                                           keepbroken)
                             klass.init()
 
-                            new_files = self._extract(klass, fid, [p.password.strip()], thread)
+                            new_files = self._extract(klass, fid, thread)
 
                         except Exception, e:
                             self.logError(basename(target), e)
+                            new_files = None
+
+                        if new_files is None:
                             success = False
                             continue
 
-                        self.logDebug("Extracted", new_files)
+                        self.logDebug("Extracted files: %s" % new_files)
                         self.setPermissions(new_files)
 
                         for file in new_files:
@@ -242,43 +267,78 @@ class ExtractArchive(Hook):
         return True if not failed else False
 
 
-    def _extract(self, plugin, fid, passwords, thread):
+    def _extract(self, plugin, fid, thread):
         pyfile = self.core.files.getFile(fid)
-        deletearchive = self.getConfig("deletearchive")
 
         pyfile.setCustomStatus(_("extracting"))
         thread.addActive(pyfile)  # keep this file until everything is done
 
         try:
-            progress = lambda x: pyfile.setProgress(x)
-            success = False
+            progress  = lambda x: pyfile.setProgress(x)
+            encrypted = False
+            passwords = self.getPasswords()
+
+            try:
+                self.logInfo(basename(plugin.file), "Verifying...")
+
+                tmp_password    = plugin.password
+                plugin.password = ""  #: Force verifying without password
+
+                plugin.verify()
+
+            except PasswordError:
+                encrypted = True
+
+            except CRCError:
+                self.logWarning(basename(plugin.file), _("Archive damaged"))
+
+                if not self.getConfig("repair"):
+                    raise CRCError
+
+                elif plugin.repair():
+                    self.logInfo(basename(plugin.file), _("Successfully repaired"))
+
+                elif not self.getConfig("keepbroken"):
+                    raise ArchiveError(_("Broken archive"))
+
+                else:
+                    self.logInfo(basename(plugin.file), _("All OK"))
+
+            plugin.password = tmp_password
+
+            if not encrypted:
+                plugin.extract(progress)
 
-            if not plugin.checkArchive():
-                plugin.extract(progress, pw)
-                success = True
             else:
                 self.logInfo(basename(plugin.file), _("Password protected"))
-                self.logDebug("Passwords: %s" % passwords if passwords else "No password provided")
 
-                for pw in set(passwords) | set(self.getPasswords()):
+                if plugin.password:
+                    passwords.insert(0, plugin.password)
+                    passwords = uniqify(self.passwords)
+                    self.logDebug("Password: %s" % plugin.password)
+                else:
+                    self.logDebug("No package password provided")
+
+                for pw in passwords:
                     try:
                         self.logDebug("Try password: %s" % pw)
-                        if plugin.checkPassword(pw):
-                            plugin.extract(progress, pw)
+
+                        if plugin.setPassword(pw):
+                            plugin.extract(progress)
                             self.addPassword(pw)
-                            success = True
                             break
+                        else:
+                            raise PasswordError
 
-                    except WrongPassword:
+                    except PasswordError:
                         self.logDebug("Password was wrong")
-
-            if not success:
-                raise Exception(_("Wrong password"))
+                else:
+                    raise PasswordError
 
             if self.core.debug:
-                self.logDebug("Would delete", ", ".join(plugin.getDeleteFiles()))
+                self.logDebug("Would delete: %s" % ", ".join(plugin.getDeleteFiles()))
 
-            if deletearchive:
+            if self.getConfig("delete"):
                 files = plugin.getDeleteFiles()
                 self.logInfo(_("Deleting %s files") % len(files))
                 for f in files:
@@ -294,12 +354,16 @@ class ExtractArchive(Hook):
 
             return extracted_files
 
-        except ArchiveError, e:
-            self.logError(basename(plugin.file), _("Archive Error"), e)
+        except PasswordError:
+            self.logError(basename(plugin.file), _("Wrong password" if passwords else "No password found"))
+            plugin.password = ""
 
         except CRCError:
             self.logError(basename(plugin.file), _("CRC Mismatch"))
 
+        except ArchiveError, e:
+            self.logError(basename(plugin.file), _("Archive Error"), e)
+
         except Exception, e:
             if self.core.debug:
                 print_exc()
@@ -307,7 +371,7 @@ class ExtractArchive(Hook):
 
         self.manager.dispatchEvent("archive_extract_failed", pyfile)
 
-        raise Exception(_("Extract failed"))
+        self.logError(basename(plugin.file), _("Extract failed"))
 
 
     @Expose
@@ -337,15 +401,13 @@ class ExtractArchive(Hook):
         """  Adds a password to saved list"""
         passwordfile = self.getConfig("passwordfile")
 
-        if pw in self.passwords:
-            self.passwords.remove(pw)
-
         self.passwords.insert(0, pw)
+        self.passwords = uniqify(self.passwords)
 
         try:
             with open(passwordfile, "wb") as f:
                 for pw in self.passwords:
-                    f.write(pw + "\n")
+                    f.write(pw + '\n')
 
         except IOError, e:
             self.logError(e)