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 import RarslaveConfig
   7 import RarslaveLogger
   8
   9 # Global Variables
  10 (TYPE_OLDRAR, TYPE_NEWRAR, TYPE_ZIP, TYPE_NOEXTRACT) = range (4)
  11 (SUCCESS, ECHECK, EEXTRACT, EDELETE) = range(4)
  12 config = RarslaveConfig.RarslaveConfig()
  13 logger = RarslaveLogger.RarslaveLogger ()
  14
  15 class RarslaveExtractor (object):
  16
  17         def __init__ (self, type):
  18                 self.type = type
  19                 self.heads = []
  20
  21         def addHead (self, dir, head):
  22                 assert os.path.isdir (dir)
  23                 # REQUIRES that the dir is valid, but not that the file is valid, so that
  24                 # we can move a file that doesn't exist yet.
  25                 # FIXME: probably CAN add this back, since we should be running this AFTER repair.
  26                 #assert os.path.isfile (os.path.join (dir, head))
  27
  28                 full_head = os.path.join (dir, head)
  29                 logger.addMessage ('Adding extraction head: %s' % full_head, RarslaveLogger.MessageType.Debug)
  30                 self.heads.append (full_head)
  31
  32         def extract (self, todir=None):
  33                 # Extract all heads of this set
  34
  35                 # Create the directory $todir if it doesn't exist
  36                 if todir != None and not os.path.isdir (todir):
  37                         logger.addMessage ('Creating directory: %s' % todir, RarslaveLogger.MessageType.Verbose)
  38                         try:
  39                                 os.makedirs (todir)
  40                         except OSError:
  41                                 logger.addMessage ('FAILED to create directory: %s' % todir, RarslaveLogger.MessageType.Fatal)
  42                                 return -EEXTRACT
  43
  44                 # Extract all heads
  45                 extraction_func = \
  46                         { TYPE_OLDRAR : self.__extract_rar,
  47                           TYPE_NEWRAR : self.__extract_rar,
  48                           TYPE_ZIP    : self.__extract_zip,
  49                           TYPE_NOEXTRACT : self.__extract_noextract }[self.type]
  50
  51                 # Call the extraction function on each head
  52                 for h in self.heads:
  53                         if todir == None:
  54                                 # Run in the head's directory
  55                                 ret = extraction_func (h, os.path.dirname (h))
  56                         else:
  57                                 ret = extraction_func (h, todir)
  58
  59                         logger.addMessage ('Extraction Function returned: %d' % ret, RarslaveLogger.MessageType.Debug)
  60
  61                         # Check error code
  62                         if ret != SUCCESS:
  63                                 logger.addMessage ('Failed extracting: %s' % h, RarslaveLogger.MessageType.Fatal)
  64                                 return -EEXTRACT
  65
  66                 return SUCCESS
  67
  68         def __extract_rar (self, file, todir):
  69                 assert os.path.isfile (file)
  70                 assert os.path.isdir (todir)
  71
  72                 RAR_CMD = config.get_value ('commands', 'unrar')
  73
  74                 cmd = '%s \"%s\"' % (RAR_CMD, file)
  75                 ret = run_command (cmd, todir)
  76
  77                 # Check error code
  78                 if ret != 0:
  79                         return -EEXTRACT
  80
  81                 return SUCCESS
  82
  83         def __extract_zip (self, file, todir):
  84                 ZIP_CMD = config.get_value ('commands', 'unzip')
  85
  86                 cmd = ZIP_CMD % (file, todir)
  87                 ret = run_command (cmd)
  88
  89                 # Check error code
  90                 if ret != 0:
  91                         return -EEXTRACT
  92
  93                 return SUCCESS
  94
  95         def __extract_noextract (self, file, todir):
  96                 # Just move this file to the $todir, since no extraction is needed
  97                 # FIXME: NOTE: mv will fail by itself if you're moving to the same dir!
  98                 NOEXTRACT_CMD = config.get_value ('commands', 'noextract')
  99
 100                 cmd = NOEXTRACT_CMD % (file, todir)
 101                 ret = run_command (cmd)
 102
 103                 # Check error code
 104                 if ret != 0:
 105                         return -EEXTRACT
 106
 107                 return SUCCESS
 108
 109
 110
 111 class RarslaveRepairer (object):
 112         # Verify (and repair) the set
 113         # Make sure it worked, otherwise clean up and return failure
 114
 115         def __init__ (self, dir, file, join=False):
 116                 self.dir  = dir  # the directory containing the par2 file
 117                 self.file = file # the par2 file
 118                 self.join = join # True if the par2 set is 001 002 ...
 119
 120                 assert os.path.isdir (dir)
 121                 assert os.path.isfile (os.path.join (dir, file))
 122
 123         def checkAndRepair (self):
 124                 # Form the command:
 125                 # par2repair -- PAR2 PAR2_EXTRA [JOIN_FILES]
 126                 PAR2_CMD = config.get_value ('commands', 'par2repair')
 127
 128                 # Get set up
 129                 basename = get_basename (self.file)
 130                 all_files = find_likely_files (basename, self.dir)
 131                 all_files.sort ()
 132                 par2_files = find_par2_files (all_files)
 133
 134                 # assemble the command
 135                 command = "%s \"%s\" " % (PAR2_CMD, self.file)
 136
 137                 for f in par2_files:
 138                         if f != self.file:
 139                                 command += "\"%s\" " % os.path.split (f)[1]
 140
 141                 if self.join:
 142                         for f in all_files:
 143                                 if f not in par2_files:
 144                                         command += "\"%s\" " % os.path.split (f)[1]
 145
 146                 # run the command
 147                 ret = run_command (command, self.dir)
 148
 149                 # check the result
 150                 if ret != 0:
 151                         logger.addMessage ('PAR2 Check / Repair failed: %s' % self.file, RarslaveLogger.MessageType.Fatal)
 152                         return -ECHECK
 153
 154                 return SUCCESS
 155
 156 def run_command (cmd, indir=None):
 157         # Runs the specified command-line in the directory given (or, in the current directory
 158         # if none is given). It returns the status code given by the application.
 159
 160         pwd = os.getcwd ()
 161
 162         if indir != None:
 163                 assert os.path.isdir (indir) # MUST be a directory!
 164                 os.chdir (indir)
 165
 166         # FIXME: re-enable this after testing
 167         print 'RUNNING (%s): %s' % (indir, cmd)
 168         return SUCCESS
 169
 170         # return os.system (cmd)
 171
 172
 173 def full_abspath (p):
 174         return os.path.abspath (os.path.expanduser (p))
 175
 176 def get_basename (name):
 177         """Strips most kinds of endings from a filename"""
 178
 179         regex = config.get_value ('regular expressions', 'basename_regex')
 180         r = re.compile (regex, re.IGNORECASE)
 181         done = False
 182
 183         while not done:
 184                 done = True
 185
 186                 if r.match (name):
 187                         g = r.match (name).groups()
 188                         name = g[0]
 189                         done = False
 190
 191         return name
 192
 193 def find_likely_files (name, dir):
 194         """Finds files which are likely to be part of the set corresponding
 195            to $name in the directory $dir"""
 196
 197         if not os.path.isdir (os.path.abspath (dir)):
 198                 raise ValueError # bad directory given
 199
 200         dir = os.path.abspath (dir)
 201         ename = re.escape (name)
 202         regex = re.compile ('^%s.*$' % (ename, ))
 203
 204         return [f for f in os.listdir (dir) if regex.match (f)]
 205
 206 def find_par2_files (files):
 207         """Find all par2 files in the list $files"""
 208
 209         PAR2_REGEX = config.get_value ('regular expressions', 'par2_regex')
 210         regex = re.compile (PAR2_REGEX, re.IGNORECASE)
 211         return [f for f in files if regex.match (f)]
 212
 213 def find_all_par2_files (dir):
 214         """Finds all par2 files in a directory"""
 215         # NOTE: does NOT return absolute paths
 216
 217         if not os.path.isdir (os.path.abspath (dir)):
 218                 raise ValueError # bad directory given
 219
 220         dir = os.path.abspath (dir)
 221         files = os.listdir (dir)
 222
 223         return find_par2_files (files)
 224
 225 def has_extension (f, ext):
 226         """Checks if f has the extension ext"""
 227
 228         if ext[0] != '.':
 229                 ext = '.' + ext
 230
 231         ext = re.escape (ext)
 232         regex = re.compile ('^.*%s$' % (ext, ), re.IGNORECASE)
 233         return regex.match (f)
 234
 235 def find_extraction_heads (dir, files):
 236         """Takes a list of possible files and finds likely heads of
 237            extraction."""
 238
 239         # NOTE: perhaps this should happen AFTER repair is
 240         # NOTE: successful. That way all files would already exist
 241
 242         # According to various sources online:
 243         # 1) pre rar-3.0: .rar .r00 .r01 ...
 244         # 2) post rar-3.0: .part01.rar .part02.rar
 245         # 3) zip all ver: .zip
 246
 247         extractor = None
 248         p2files = find_par2_files (files)
 249
 250         # Old RAR type, find all files ending in .rar
 251         if is_oldrar (files):
 252                 extractor = RarslaveExtractor (TYPE_OLDRAR)
 253                 regex = re.compile ('^.*\.rar$', re.IGNORECASE)
 254                 for f in files:
 255                         if regex.match (f):
 256                                 extractor.addHead (dir, f)
 257
 258         if is_newrar (files):
 259                 extractor = RarslaveExtractor (TYPE_NEWRAR)
 260                 regex = re.compile ('^.*\.part01.rar$', re.IGNORECASE)
 261                 for f in files:
 262                         if regex.match (f):
 263                                 extractor.addHead (dir, f)
 264
 265         if is_zip (files):
 266                 extractor = RarslaveExtractor (TYPE_ZIP)
 267                 regex = re.compile ('^.*\.zip$', re.IGNORECASE)
 268                 for f in files:
 269                         if regex.match (f):
 270                                 extractor.addHead (dir, f)
 271
 272         if is_noextract (files):
 273                 # Use the Par2 Parser (from cfv) here to find out what files are protected.
 274                 # Since these are not being extracted, they will be mv'd to another directory
 275                 # later.
 276                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
 277
 278                 for f in p2files:
 279                         done = False
 280                         try:
 281                                 prot_files = par2parser.get_protected_files (dir, f)
 282                                 done = True
 283                         except: #FIXME: add the actual exceptions
 284                                 logger.addMessage ('Error parsing PAR2 file: %s', f)
 285                                 continue
 286
 287                         if done:
 288                                 break
 289
 290                 if done:
 291                         for f in prot_files:
 292                                 extractor.addHead (dir, f)
 293                 else:
 294                         logger.addMessage ('Error parsing all PAR2 files in this set ...')
 295
 296         # Make sure we found the type
 297         if extractor == None:
 298                 logger.addMessage ('Not able to find an extractor for this type of set: %s' % p2files[0],
 299                                 RarslaveLogger.MessageType.Fatal)
 300
 301                 # No-heads here, but it's better than failing completely
 302                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
 303
 304         return extractor
 305
 306 def is_oldrar (files):
 307         for f in files:
 308                 if has_extension (f, '.r00'):
 309                         return True
 310
 311         return False
 312
 313 def is_newrar (files):
 314         for f in files:
 315                 if has_extension (f, '.part01.rar'):
 316                         return True
 317
 318         return False
 319
 320 def is_zip (files):
 321         for f in files:
 322                 if has_extension (f, '.zip'):
 323                         return True
 324
 325         return False
 326
 327 def is_noextract (files):
 328         # Type that needs no extraction.
 329         # TODO: Add others ???
 330         for f in files:
 331                 if has_extension (f, '.001'):
 332                         return True
 333
 334         return False
 335
 336 def find_deleteable_files (files):
 337         # Deleteable types regex should come from the config
 338         dfiles = []
 339         DELETE_REGEX = config.get_value ('regular expressions', 'delete_regex')
 340         dregex = re.compile (DELETE_REGEX, re.IGNORECASE)
 341
 342         return [f for f in files if dregex.match (f)]
 343
 344 def printlist (li):
 345         for f in li:
 346                 print f
 347
 348 class PAR2Set (object):
 349
 350         dir = None
 351         file = None
 352         likely_files = []
 353
 354         def __init__ (self, dir, file):
 355                 assert os.path.isdir (dir)
 356                 assert os.path.isfile (os.path.join (dir, file))
 357
 358                 self.dir = dir
 359                 self.file = file
 360
 361                 basename = get_basename (file)
 362                 self.likely_files = find_likely_files (basename, dir)
 363
 364         def __list_eq (self, l1, l2):
 365
 366                 if len(l1) != len(l2):
 367                         return False
 368
 369                 for e in l1:
 370                         if e not in l2:
 371                                 return False
 372
 373                 return True
 374
 375         def __eq__ (self, rhs):
 376                 return self.__list_eq (self.likely_files, rhs.likely_files)
 377
 378         def run_all (self):
 379                 par2files = find_par2_files (self.likely_files)
 380                 par2head = par2files[0]
 381
 382                 join = is_noextract (self.likely_files)
 383
 384                 # Repair Stage
 385                 repairer = RarslaveRepairer (self.dir, par2head, join)
 386                 ret = repairer.checkAndRepair ()
 387
 388                 if ret != SUCCESS:
 389                         logger.addMessage ('Repair stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
 390                         return -ECHECK
 391
 392                 # Extraction Stage
 393                 EXTRACT_DIR = config.get_value ('directories', 'extract_directory')
 394                 extractor = find_extraction_heads (self.dir, self.likely_files)
 395                 ret = extractor.extract (EXTRACT_DIR)
 396
 397                 if ret != SUCCESS:
 398                         logger.addMessage ('Extraction stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
 399                         return -EEXTRACT
 400
 401                 # Deletion Stage
 402                 DELETE_INTERACTIVE = config.get_value ('options', 'interactive')
 403                 deleteable_files = find_deleteable_files (self.likely_files)
 404                 ret = delete_list (deleteable_files, DELETE_INTERACTIVE)
 405
 406                 if ret != SUCCESS:
 407                         logger.addMessage ('Deletion stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
 408                         return -EDELETE
 409
 410                 logger.addMessage ('Successfully completed: %s' % par2head)
 411                 return SUCCESS
 412
 413 def delete_list (files, interactive=False):
 414         # Delete a list of files
 415
 416         done = False
 417         valid_y = ['Y', 'YES']
 418         valid_n = ['N', 'NO']
 419
 420         if interactive:
 421                 while not done:
 422                         print 'Do you want to delete the following?:'
 423                         s = raw_input ('Delete [y/N]: ').upper()
 424
 425                         if s in valid_y + valid_n:
 426                                 done = True
 427
 428                 if s in valid_n:
 429                         return SUCCESS
 430
 431         for f in files:
 432                 # FIXME: re-enable this in production
 433                 # os.remove (f)
 434                 print 'rm \"%s\"' % f
 435
 436         return SUCCESS
 437
 438
 439 def generate_all_parsets (dir):
 440         # Generate all parsets in the given directory.
 441
 442         assert os.path.isdir (dir) # Directory MUST be valid
 443
 444         parsets = []
 445         p2files = find_all_par2_files (dir)
 446
 447         for f in p2files:
 448                 p = PAR2Set (dir, f)
 449                 if p not in parsets:
 450                         parsets.append (p)
 451
 452         return parsets
 453
 454 def main ():
 455         TOPDIR = os.path.abspath ('test_material')
 456
 457         for (dir, subdirs, files) in os.walk (TOPDIR):
 458                 parsets = generate_all_parsets (dir)
 459                 for p in parsets:
 460                         p.run_all ()
 461
 462         print '\nRARSLAVE STATUS\n'
 463
 464         # Used in '--quiet' mode
 465         if logger.hasFatalMessages ():
 466                 print '\nFatal Messages\n' + '=' * 80
 467                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
 468
 469         # Used in no options mode
 470         if logger.hasNormalMessages ():
 471                 print '\nNormal Messages\n' + '=' * 80
 472                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
 473
 474         # Used in --verbose mode
 475         if logger.hasVerboseMessages ():
 476                 print '\nVerbose Messages\n' + '=' * 80
 477                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
 478
 479         # Used in --debug mode
 480         if logger.hasDebugMessages ():
 481                 print '\nDebug Messages\n' + '=' * 80
 482                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
 483
 484         print '\n\nALL MESSAGES:'
 485         logger.printAllMessages ()
 486
 487 if __name__ == '__main__':
 488         main ()
 489