update uloz.to, cloudnator.com, checksum plugin

1 files changed, 41 insertions, 15 deletions

diff --git a/module/plugins/hooks/Checksum.py b/module/plugins/hooks/Checksum.py
index 89e8ec762..0c2751e7a 100644
--- a/module/plugins/hooks/Checksum.py
+++ b/module/plugins/hooks/Checksum.py
@@ -18,6 +18,7 @@
 """
 from __future__ import with_statement
 import hashlib, zlib
+from os.path import getsize, isfile
 
 from module.utils import save_join, fs_encode
 from module.plugins.Hook import Hook
@@ -48,13 +49,20 @@ def computeChecksum(local_file, algorithm):
 
 class Checksum(Hook):
     __name__ = "Checksum"
-    __version__ = "0.03"
-    __description__ = "Check downloaded file hash"
+    __version__ = "0.04"
+    __description__ = "Verify downloaded file size and checksum (enable in general preferences)"
     __config__ = [("activated", "bool", "Activated", True),
                   ("action", "fail;retry;nothing", "What to do if check fails?", "retry"),
                   ("max_tries", "int", "Number of retries", 2)]
     __author_name__ = ("zoidberg")
     __author_mail__ = ("zoidberg@mujmail.cz")
+    
+    def setup(self):    
+        self.algorithms = sorted(getattr(hashlib, "algorithms", ("md5", "sha1", "sha224", "sha256", "sha384", "sha512")), reverse = True)
+        self.algorithms.append(["crc32", "adler32"])
+        
+        if not self.config['general']['checksum']:
+            self.logInfo("Checksum validation is disabled in general configuration")                                  
              
     def downloadFinished(self, pyfile):
         """ 
@@ -62,25 +70,43 @@ class Checksum(Hook):
         pyfile.plugin.check_data should be a dictionary which can contain:
         a) if known, the exact filesize in bytes (e.g. "size": 123456789)
         b) hexadecimal hash string with algorithm name as key (e.g. "md5": "d76505d0869f9f928a17d42d66326307")    
-        """
+        """        
+        
         if hasattr(pyfile.plugin, "check_data") and (isinstance(pyfile.plugin.check_data, dict)):
+            data = pyfile.plugin.check_data        
+        elif hasattr(pyfile.plugin, "api_data") and (isinstance(pyfile.plugin.api_data, dict)):
+            data = pyfile.plugin.api_data
+        else:
+            return                    
+        
+        download_folder = self.config['general']['download_folder']
+        local_file = fs_encode(save_join(download_folder, pyfile.package().folder, pyfile.name))
+        
+        if not isfile(local_file):
+            self.checkFailed(pyfile, "File does not exist")  
+        
+        # validate file size
+        if "size" in data:
+            api_size = int(data['size'])
+            file_size = getsize(local_file)
+            if api_size != file_size:
+                self.logWarning("File %s has incorrect size: %d B (%d expected)" % (pyfile.name, file_size, api_size))
+                self.checkFailed(pyfile, "Incorrect file size")
+                
+        # validate checksum
+        if self.config['general']['checksum']:                                                       
+            if "checksum" in data:
+                data['md5'] = data['checksum']
             
-            download_folder = self.config['general']['download_folder']
-            local_file = fs_encode(save_join(download_folder, pyfile.package().folder, pyfile.name))
-            
-            for key, value in sorted(pyfile.plugin.check_data.items(), reverse = True):                          
-                if key == "size":
-                    if value and value != pyfile.size:
-                        self.logWarning("File %s has incorrect size: %d B (%d expected)" % (pyfile.size, value))
-                        self.checkFailed(pyfile, "Incorrect file size")
-                else:
+            for key in self.algorithms:
+                if key in data: 
                     checksum = computeChecksum(local_file, key.replace("-","").lower())                    
                     if checksum:
-                        if checksum == value:
-                            self.logInfo('File integrity of "%s" verified by %s checksum (%s).' % (pyfile.name, key.upper() , checksum))
+                        if checksum == data[key]:
+                            self.logInfo('File integrity of "%s" verified by %s checksum (%s).' % (pyfile.name, key.upper(), checksum))
                             return
                         else:
-                            self.logWarning("%s checksum for file %s does not match (%s != %s)" % (key.upper(), pyfile.name, checksum, value))    
+                            self.logWarning("%s checksum for file %s does not match (%s != %s)" % (key.upper(), pyfile.name, checksum, data[key]))    
                             self.checkFailed(pyfile, "Checksums do not match")
                     else:
                         self.logWarning("Unsupported hashing algorithm: %s" % key.upper())