rarslave.py

   1 #!/usr/bin/env python
   2 # vim: set ts=4 sts=4 sw=4 textwidth=112 :
   3
   4 import re, os, sys
   5 import par2parser
   6
   7 # Global Variables
   8 (TYPE_OLDRAR, TYPE_NEWRAR, TYPE_ZIP, TYPE_NOEXTRACT) = range (4)
   9 (ECHECK, EEXTRACT, EDELETE) = range(1,4)
  10
  11 class RarslaveExtractor (object):
  12
  13         def __init__ (self, type):
  14                 self.type = type
  15                 self.heads = []
  16
  17         def addHead (self, dir, head):
  18                 assert os.path.isdir (dir)
  19                 # REQUIRES that the dir is valid, but not that the file is valid, so that
  20                 # we can move a file that doesn't exist yet.
  21                 # FIXME: probably CAN add this back, since we should be running this AFTER repair.
  22                 #assert os.path.isfile (os.path.join (dir, head))
  23
  24                 self.heads.append (os.path.join (dir, head))
  25
  26         def extract (self, todir):
  27                 # Extract all heads of this set
  28
  29                 # Create the directory $todir if it doesn't exist
  30                 if not os.path.isdir (todir):
  31                         # TODO: LOGGER
  32                         try:
  33                                 os.makedirs (todir)
  34                         except OSError:
  35                                 # TODO: LOGGER
  36                                 return -EEXTRACT
  37
  38                 # Extract all heads
  39                 extraction_func = \
  40                         { TYPE_OLDRAR : self.__extract_rar,
  41                           TYPE_NEWRAR : self.__extract_rar,
  42                           TYPE_ZIP    : self.__extract_zip,
  43                           TYPE_NOEXTRACT : self.__extract_noextract }[self.type]
  44
  45                 # Call the extraction function on each head
  46                 for h in self.heads:
  47                         extraction_func (h, todir)
  48
  49         def __extract_rar (self, file, todir):
  50                 assert os.path.isfile (file)
  51                 assert os.path.isdir (todir)
  52
  53                 RAR_CMD = 'unrar x -o+ -- '
  54
  55                 cmd = '%s \"%s\"' % (RAR_CMD, file)
  56                 ret = run_command (cmd, todir)
  57
  58                 # Check error code
  59                 if ret != 0:
  60                         return -EEXTRACT
  61
  62         def __extract_zip (self, file, todir):
  63                 ZIP_CMD = 'unzip \"%s\" -d \"%s\"'
  64
  65                 cmd = ZIP_CMD % (file, todir)
  66                 ret = run_command (cmd)
  67
  68                 # Check error code
  69                 if ret != 0:
  70                         return -EEXTRACT
  71
  72         def __extract_noextract (self, file, todir):
  73                 # Just move this file to the $todir, since no extraction is needed
  74                 # FIXME: NOTE: mv will fail by itself if you're moving to the same dir!
  75                 cmd = 'mv \"%s\" \"%s\"' % (file, todir)
  76                 ret = run_command (cmd)
  77
  78                 # Check error code
  79                 if ret != 0:
  80                         return -EEXTRACT
  81
  82
  83
  84 class RarslaveRepairer (object):
  85         # Verify (and repair) the set
  86         # Make sure it worked, otherwise clean up and return failure
  87
  88         def __init__ (self, dir, file, join=False):
  89                 self.dir  = dir  # the directory containing the par2 file
  90                 self.file = file # the par2 file
  91                 self.join = join # True if the par2 set is 001 002 ...
  92
  93                 assert os.path.isdir (dir)
  94                 assert os.path.isfile (os.path.join (dir, file))
  95
  96         def checkAndRepair (self):
  97                 # Form the command:
  98                 # par2repair -- PAR2 PAR2_EXTRA [JOIN_FILES]
  99                 PAR2_CMD = 'par2repair -- '
 100
 101                 # Get set up
 102                 basename = get_basename (self.file)
 103                 all_files = find_likely_files (basename, self.dir)
 104                 all_files.sort ()
 105                 par2_files = find_par2_files (all_files)
 106
 107                 # assemble the command
 108                 command = "%s \"%s\" " % (PAR2_CMD, self.file)
 109
 110                 for f in par2_files:
 111                         if f != self.file:
 112                                 command += "\"%s\" " % get_filename(f)
 113
 114                 if self.join:
 115                         for f in all_files:
 116                                 if f not in par2_files:
 117                                         command += "\"%s\" " % get_filename(f)
 118
 119                 # run the command
 120                 ret = run_command (command, self.dir)
 121
 122                 # check the result
 123                 if ret != 0:
 124                         # TODO: logger
 125                         print 'error during checkAndRepair()'
 126                         return -ECHECK
 127
 128 def run_command (cmd, indir=None):
 129         # Runs the specified command-line in the directory given (or, in the current directory
 130         # if none is given). It returns the status code given by the application.
 131
 132         pwd = os.getcwd ()
 133
 134         if indir != None:
 135                 assert os.path.isdir (indir) # MUST be a directory!
 136                 os.chdir (pwd)
 137
 138         # FIXME: re-enable this after testing
 139         print 'RUNNING (%s): %s' % (indir, cmd)
 140         return 0
 141
 142         # return os.system (cmd)
 143
 144
 145 def full_abspath (p):
 146         return os.path.abspath (os.path.expanduser (p))
 147
 148 def get_filename (f):
 149         # TODO: I don't think that we should enforce this...
 150         # TODO: ... because I think we should be able to get the filename, regardless
 151         # TODO: of whether this is a legit filename RIGHT NOW or not.
 152         # assert os.path.isfile (f)
 153         return os.path.split (f)[1]
 154
 155 def get_basename (name):
 156         """Strips most kinds of endings from a filename"""
 157
 158         regex = '^(.+)\.(par2|vol\d+\+\d+|\d\d\d|part\d+|rar|zip|avi|mp4|mkv|ogm)$'
 159         r = re.compile (regex, re.IGNORECASE)
 160         done = False
 161
 162         while not done:
 163                 done = True
 164
 165                 if r.match (name):
 166                         g = r.match (name).groups()
 167                         name = g[0]
 168                         done = False
 169
 170         return name
 171
 172 def find_likely_files (name, dir):
 173         """Finds files which are likely to be part of the set corresponding
 174            to $name in the directory $dir"""
 175
 176         if not os.path.isdir (os.path.abspath (dir)):
 177                 raise ValueError # bad directory given
 178
 179         dir = os.path.abspath (dir)
 180         ename = re.escape (name)
 181         regex = re.compile ('^%s.*$' % (ename, ))
 182
 183         return [f for f in os.listdir (dir) if regex.match (f)]
 184
 185 def find_par2_files (files):
 186         """Find all par2 files in the list $files"""
 187
 188         regex = re.compile ('^.*\.par2$', re.IGNORECASE)
 189         return [f for f in files if regex.match (f)]
 190
 191 def find_all_par2_files (dir):
 192         """Finds all par2 files in a directory"""
 193         # NOTE: does NOT return absolute paths
 194
 195         if not os.path.isdir (os.path.abspath (dir)):
 196                 raise ValueError # bad directory given
 197
 198         dir = os.path.abspath (dir)
 199         files = os.listdir (dir)
 200
 201         return find_par2_files (files)
 202
 203 def has_extension (f, ext):
 204         """Checks if f has the extension ext"""
 205
 206         if ext[0] != '.':
 207                 ext = '.' + ext
 208
 209         ext = re.escape (ext)
 210         regex = re.compile ('^.*%s$' % (ext, ), re.IGNORECASE)
 211         return regex.match (f)
 212
 213 def find_extraction_heads (dir, files):
 214         """Takes a list of possible files and finds likely heads of
 215            extraction."""
 216
 217         # NOTE: perhaps this should happen AFTER repair is
 218         # NOTE: successful. That way all files would already exist
 219
 220         # According to various sources online:
 221         # 1) pre rar-3.0: .rar .r00 .r01 ...
 222         # 2) post rar-3.0: .part01.rar .part02.rar
 223         # 3) zip all ver: .zip
 224
 225         extractor = None
 226         p2files = find_par2_files (files)
 227
 228         # Old RAR type, find all files ending in .rar
 229         if is_oldrar (files):
 230                 extractor = RarslaveExtractor (TYPE_OLDRAR)
 231                 regex = re.compile ('^.*\.rar$', re.IGNORECASE)
 232                 for f in files:
 233                         if regex.match (f):
 234                                 extractor.addHead (dir, f)
 235
 236         if is_newrar (files):
 237                 extractor = RarslaveExtractor (TYPE_NEWRAR)
 238                 regex = re.compile ('^.*\.part01.rar$', re.IGNORECASE)
 239                 for f in files:
 240                         if regex.match (f):
 241                                 extractor.addHead (dir, f)
 242
 243         if is_zip (files):
 244                 extractor = RarslaveExtractor (TYPE_ZIP)
 245                 regex = re.compile ('^.*\.zip$', re.IGNORECASE)
 246                 for f in files:
 247                         if regex.match (f):
 248                                 extractor.addHead (dir, f)
 249
 250         if is_noextract (files):
 251                 # Use the Par2 Parser (from cfv) here to find out what files are protected.
 252                 # Since these are not being extracted, they will be mv'd to another directory
 253                 # later.
 254                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
 255
 256                 for f in p2files:
 257                         done = False
 258                         try:
 259                                 prot_files = par2parser.get_protected_files (dir, f)
 260                                 done = True
 261                         except: #FIXME: add the actual exceptions
 262                                 print 'ERROR PARSING P2FILE ...', f
 263                                 continue
 264
 265                         if done:
 266                                 break
 267
 268                 if done:
 269                         for f in prot_files:
 270                                 extractor.addHead (dir, f)
 271                 else:
 272                         print 'BADNESS'
 273
 274         # Make sure we found the type
 275         assert extractor != None
 276
 277         return extractor
 278
 279 def is_oldrar (files):
 280         for f in files:
 281                 if has_extension (f, '.r00'):
 282                         return True
 283
 284 def is_newrar (files):
 285         for f in files:
 286                 if has_extension (f, '.part01.rar'):
 287                         return True
 288
 289 def is_zip (files):
 290         for f in files:
 291                 if has_extension (f, '.zip'):
 292                         return True
 293
 294 def is_noextract (files):
 295         # Type that needs no extraction.
 296         # TODO: Add others ???
 297         for f in files:
 298                 if has_extension (f, '.001'):
 299                         return True
 300
 301 def find_deleteable_files (files):
 302         # Deleteable types regex should come from the config
 303         dfiles = []
 304         dregex = re.compile ('^.*\.(par2|\d|\d\d\d|rar|r\d\d|zip)$', re.IGNORECASE)
 305
 306         return [f for f in files if dregex.match (f)]
 307
 308 def printlist (li):
 309         for f in li:
 310                 print f
 311
 312 class PAR2Set (object):
 313
 314         dir = None
 315         file = None
 316         likely_files = []
 317
 318         def __init__ (self, dir, file):
 319                 assert os.path.isdir (dir)
 320                 assert os.path.isfile (os.path.join (dir, file))
 321
 322                 self.dir = dir
 323                 self.file = file
 324
 325                 basename = get_basename (file)
 326                 self.likely_files = find_likely_files (basename, dir)
 327
 328         def __list_eq (self, l1, l2):
 329
 330                 if len(l1) != len(l2):
 331                         return False
 332
 333                 for e in l1:
 334                         if e not in l2:
 335                                 return False
 336
 337                 return True
 338
 339         def __eq__ (self, rhs):
 340                 return self.__list_eq (self.likely_files, rhs.likely_files)
 341
 342         def run_all (self):
 343                 par2files = find_par2_files (self.likely_files)
 344                 par2head = par2files[0]
 345
 346                 join = is_noextract (self.likely_files)
 347
 348                 # Repair Stage
 349                 repairer = RarslaveRepairer (self.dir, par2head, join)
 350                 ret = repairer.checkAndRepair () # FIXME: Check return value
 351
 352                 if ret: # FAILURE
 353                         return -ECHECK
 354
 355                 # Extraction Stage
 356                 extractor = find_extraction_heads (self.dir, self.likely_files)
 357                 ret = extractor.extract ('extract_dir') # FIXME: Get it from the config
 358
 359                 if ret: # FAILURE
 360                         return -EEXTRACT
 361
 362                 # Deletion Stage
 363                 deleteable_files = find_deleteable_files (self.likely_files)
 364                 ret = delete_list (deleteable_files)
 365
 366                 if ret: # FAILURE
 367                         return -EDELETE
 368
 369                 return 0
 370
 371 def delete_list (files, interactive=False):
 372         # Delete a list of files
 373         # TODO: Add the ability to confirm deletion, like in the original rarslave
 374
 375         if interactive:
 376                 # TODO: prompt here
 377                 # prompt -> OK_TO_DELETE -> do nothing, fall through
 378                 # prompt -> NOT_OK -> return immediately
 379                 pass
 380
 381         for f in files:
 382                 # FIXME: re-enable this in production
 383                 # os.remove (f)
 384                 print 'rm', f
 385
 386         return 0
 387
 388
 389 def generate_all_parsets (dir):
 390         # Generate all parsets in the given directory.
 391
 392         assert os.path.isdir (dir) # Directory MUST be valid
 393
 394         parsets = []
 395         p2files = find_all_par2_files (dir)
 396
 397         for f in p2files:
 398                 p = PAR2Set (dir, f)
 399                 if p not in parsets:
 400                         parsets.append (p)
 401
 402         return parsets
 403
 404 def main ():
 405         TOPDIR = os.path.abspath ('test_material')
 406
 407         for (dir, subdirs, files) in os.walk (TOPDIR):
 408                 print 'DEBUG: IN DIRECTORY:', dir
 409                 parsets = generate_all_parsets (dir)
 410                 for p in parsets:
 411                         p.run_all ()
 412
 413 if __name__ == '__main__':
 414         main ()