rarslave.py

   1 #!/usr/bin/env python
   2 # vim: set ts=4 sts=4 sw=4 textwidth=112 :
   3
   4 VERSION="2.0.0"
   5 PROGRAM="rarslave2"
   6
   7 import re, os, sys, optparse
   8 import Par2Parser
   9 import RarslaveConfig
  10 import RarslaveLogger
  11
  12 # Global Variables
  13 (TYPE_OLDRAR, TYPE_NEWRAR, TYPE_ZIP, TYPE_NOEXTRACT, TYPE_UNKNOWN) = range (5)
  14 (SUCCESS, ECHECK, EEXTRACT, EDELETE) = range(4)
  15 config = RarslaveConfig.RarslaveConfig()
  16 logger = RarslaveLogger.RarslaveLogger ()
  17
  18 # Global options to be set / used later.
  19 options = None
  20
  21 class RarslaveExtractor (object):
  22
  23         type = None
  24         heads = []
  25
  26         def __init__ (self, dir, p2files, name_files, prot_files):
  27
  28                 self.dir = dir
  29                 self.p2files = p2files
  30                 self.name_matched_files = name_files
  31                 self.prot_matched_files = prot_files
  32
  33                 # Find the type
  34                 self.type = self.__find_type ()
  35
  36                 logger.addMessage ('Detected set of type: %s' % self, RarslaveLogger.MessageType.Debug)
  37
  38                 # Find the heads
  39                 self.heads = self.__find_heads ()
  40
  41                 for h in self.heads:
  42                         logger.addMessage ('Adding extraction head: %s' % h, RarslaveLogger.MessageType.Debug)
  43
  44         def __repr__ (self):
  45                 return \
  46                         {       TYPE_OLDRAR : 'Old RAR',
  47                                 TYPE_NEWRAR : 'New RAR',
  48                                 TYPE_ZIP :    'Zip',
  49                                 TYPE_NOEXTRACT : 'No Extract',
  50                                 TYPE_UNKNOWN : 'Unknown' } [self.type]
  51
  52         def __find_type (self):
  53
  54                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
  55
  56                 if   self.is_oldrar (all_files):
  57                         return TYPE_OLDRAR
  58                 elif self.is_newrar (all_files):
  59                         return TYPE_NEWRAR
  60                 elif self.is_zip (all_files):
  61                         return TYPE_ZIP
  62                 elif self.is_noextract (all_files):
  63                         return TYPE_NOEXTRACT
  64
  65                 return TYPE_UNKNOWN
  66
  67         def __generic_find_heads (self, regex, ignorecase=True):
  68
  69                 heads = []
  70
  71                 if ignorecase:
  72                         cregex = re.compile (regex, re.IGNORECASE)
  73                 else:
  74                         cregex = re.compile (regex)
  75
  76                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
  77
  78                 for f in all_files:
  79                         if cregex.match (f):
  80                                 heads.append (f)
  81
  82                 return heads
  83
  84         def __find_heads (self):
  85
  86                 if self.type == TYPE_OLDRAR:
  87                         return self.__generic_find_heads ('^.*\.rar$')
  88                 elif self.type == TYPE_NEWRAR:
  89                         return self.__generic_find_heads ('^.*\.part0*1\.rar$')
  90                 elif self.type == TYPE_ZIP:
  91                         return self.__generic_find_heads ('^.*\.zip$')
  92                 elif self.type == TYPE_NOEXTRACT:
  93                         return self.prot_matched_files
  94
  95                 return []
  96
  97         def __create_directory (self, dir):
  98                 if dir == None:
  99                         return SUCCESS
 100
 101                 if os.path.isdir (dir):
 102                         return SUCCESS
 103
 104                 try:
 105                         os.makedirs (dir)
 106                         logger.addMessage ('Created directory: %s' % dir, RarslaveLogger.MessageType.Verbose)
 107                 except OSError:
 108                         logger.addMessage ('FAILED to create directory: %s' % dir, RarslaveLogger.MessageType.Fatal)
 109                         return -EEXTRACT
 110
 111                 return SUCCESS
 112
 113         def runExtract (self, todir=None):
 114                 # Extract all heads of this set
 115
 116                 # Extract to the head's dir if we don't care where to extract
 117                 if todir == None:
 118                         todir = self.dir
 119
 120                 # Create the directory $todir if it doesn't exist
 121                 ret = self.__create_directory (todir)
 122
 123                 if ret != SUCCESS:
 124                         return -EEXTRACT
 125
 126                 # Extract all heads
 127                 extraction_func = \
 128                         { TYPE_OLDRAR : self.__extract_rar,
 129                           TYPE_NEWRAR : self.__extract_rar,
 130                           TYPE_ZIP    : self.__extract_zip,
 131                           TYPE_NOEXTRACT : self.__extract_noextract,
 132                           TYPE_UNKNOWN : self.__extract_unknown }[self.type]
 133
 134                 # Call the extraction function on each head
 135                 for h in self.heads:
 136                         full_head = full_abspath (h)
 137                         ret = extraction_func (full_head, todir)
 138                         logger.addMessage ('Extraction Function returned: %d' % ret, RarslaveLogger.MessageType.Debug)
 139
 140                         # Check error code
 141                         if ret != SUCCESS:
 142                                 logger.addMessage ('Failed extracting: %s' % h, RarslaveLogger.MessageType.Fatal)
 143                                 return -EEXTRACT
 144
 145                 return SUCCESS
 146
 147         def __extract_rar (self, file, todir):
 148                 assert os.path.isfile (file)
 149                 assert os.path.isdir (todir)
 150
 151                 RAR_CMD = config.get_value ('commands', 'unrar')
 152
 153                 cmd = '%s \"%s\"' % (RAR_CMD, file)
 154                 ret = run_command (cmd, todir)
 155
 156                 # Check error code
 157                 if ret != 0:
 158                         return -EEXTRACT
 159
 160                 return SUCCESS
 161
 162         def __extract_zip (self, file, todir):
 163                 ZIP_CMD = config.get_value ('commands', 'unzip')
 164
 165                 cmd = ZIP_CMD % (file, todir)
 166                 ret = run_command (cmd)
 167
 168                 # Check error code
 169                 if ret != 0:
 170                         return -EEXTRACT
 171
 172                 return SUCCESS
 173
 174         def __extract_noextract (self, file, todir):
 175                 # Just move this file to the $todir, since no extraction is needed
 176                 # FIXME: NOTE: mv will fail by itself if you're moving to the same dir!
 177                 NOEXTRACT_CMD = config.get_value ('commands', 'noextract')
 178
 179                 # Make sure that both files are not the same file. If they are, don't run at all.
 180                 if os.path.samefile (file, os.path.join (todir, file)):
 181                         return SUCCESS
 182
 183                 cmd = NOEXTRACT_CMD % (file, todir)
 184                 ret = run_command (cmd)
 185
 186                 # Check error code
 187                 if ret != 0:
 188                         return -EEXTRACT
 189
 190                 return SUCCESS
 191
 192         def __extract_unknown (self, file, todir):
 193                 return SUCCESS
 194
 195         def __generic_matcher (self, files, regex, nocase=False):
 196                 """Run the regex over the files, and see if one matches or not.
 197                 NOTE: this does not return the matches, just if a match occurred."""
 198
 199                 if nocase:
 200                         cregex = re.compile (regex, re.IGNORECASE)
 201                 else:
 202                         cregex = re.compile (regex)
 203
 204                 for f in files:
 205                         if cregex.match (f):
 206                                 return True
 207
 208                 return False
 209
 210         def is_oldrar (self, files):
 211                 return self.__generic_matcher (files, '^.*\.r00$')
 212
 213         def is_newrar (self, files):
 214                 return self.__generic_matcher (files, '^.*\.part0*1\.rar$')
 215
 216         def is_zip (self, files):
 217                 return self.__generic_matcher (files, '^.*\.zip$')
 218
 219         def is_noextract (self, files):
 220                 # Type that needs no extraction.
 221                 # TODO: Add others ???
 222                 return self.__generic_matcher (files, '^.*\.001$')
 223
 224 class PAR2Set (object):
 225
 226         dir = None
 227         p2file = None # The starting par2
 228         basename = None # The p2file's basename
 229         all_p2files = []
 230         name_matched_files = [] # Files that match by basename of the p2file
 231         prot_matched_files = [] # Files that match by being protected members
 232
 233         def __init__ (self, dir, p2file):
 234                 assert os.path.isdir (dir)
 235                 assert os.path.isfile (os.path.join (dir, p2file))
 236
 237                 self.dir = dir
 238                 self.p2file = p2file
 239                 self.basename = self.__get_basename (p2file)
 240
 241                 # Find files that match by name only
 242                 self.name_matched_files = self.__find_name_matches (self.dir, self.basename)
 243
 244                 # Find all par2 files for this set using name matches
 245                 self.all_p2files = find_par2_files (self.name_matched_files)
 246
 247                 # Try to get the protected files for this set
 248                 self.prot_matched_files = self.__parse_all_par2 ()
 249
 250         def __list_eq (self, l1, l2):
 251
 252                 if len(l1) != len(l2):
 253                         return False
 254
 255                 for e in l1:
 256                         if e not in l2:
 257                                 return False
 258
 259                 return True
 260
 261         def __eq__ (self, rhs):
 262                 return (self.dir == rhs.dir) and (self.basename == rhs.basename) and \
 263                                 self.__list_eq (self.name_matched_files, rhs.name_matched_files) and \
 264                                 self.__list_eq (self.prot_matched_files, rhs.prot_matched_files)
 265
 266         def __get_basename (self, name):
 267                 """Strips most kinds of endings from a filename"""
 268
 269                 regex = config.get_value ('regular expressions', 'basename_regex')
 270                 r = re.compile (regex, re.IGNORECASE)
 271                 done = False
 272
 273                 while not done:
 274                         done = True
 275
 276                         if r.match (name):
 277                                 g = r.match (name).groups()
 278                                 name = g[0]
 279                                 done = False
 280
 281                 return name
 282
 283         def __parse_all_par2 (self):
 284                 """Searches though self.all_p2files and tries to parse at least one of them"""
 285                 done = False
 286                 files = []
 287
 288                 for f in self.all_p2files:
 289
 290                         # Exit early if we've found a good file
 291                         if done:
 292                                 break
 293
 294                         try:
 295                                 files = Par2Parser.get_protected_files (self.dir, f)
 296                                 done = True
 297                         except (EnvironmentError, OSError, OverflowError):
 298                                 logger.addMessage ('Corrupt PAR2 file: %s' % f, RarslaveLogger.MessageType.Fatal)
 299
 300                 # Now that we're out of the loop, check if we really finished
 301                 if not done:
 302                         logger.addMessage ('All PAR2 files corrupt for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 303
 304                 # Return whatever we've got, empty or not
 305                 return files
 306
 307         def __find_name_matches (self, dir, basename):
 308                 """Finds files which are likely to be part of the set corresponding
 309                    to $name in the directory $dir"""
 310
 311                 assert os.path.isdir (dir)
 312
 313                 ename = re.escape (basename)
 314                 regex = re.compile ('^%s.*$' % (ename, ))
 315
 316                 return [f for f in os.listdir (dir) if regex.match (f)]
 317
 318         def __update_name_matches (self):
 319                 """Updates the self.name_matched_files variable with the most current information.
 320                    This should be called after the directory contents are likely to change."""
 321
 322                 self.name_matched_files = self.__find_name_matches (self.dir, self.basename)
 323
 324         def __should_be_joined (self, files):
 325                 regex = re.compile ('^.*\.001$', re.IGNORECASE)
 326                 for f in files:
 327                         if regex.match (f):
 328                                 return True
 329
 330         def runCheckAndRepair (self):
 331                 PAR2_CMD = config.get_value ('commands', 'par2repair')
 332
 333                 # Get set up
 334                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 335                 join = self.__should_be_joined (all_files)
 336
 337                 # assemble the command
 338                 # par2repair -- PAR2 PAR2_EXTRA [JOIN_FILES]
 339                 command = "%s \"%s\" " % (PAR2_CMD, self.p2file)
 340
 341                 for f in self.all_p2files:
 342                         if f != self.p2file:
 343                                 command += "\"%s\" " % os.path.split (f)[1]
 344
 345                 if join:
 346                         for f in all_files:
 347                                 if f not in self.p2files:
 348                                         command += "\"%s\" " % os.path.split (f)[1]
 349
 350                 # run the command
 351                 ret = run_command (command, self.dir)
 352
 353                 # check the result
 354                 if ret != 0:
 355                         logger.addMessage ('PAR2 Check / Repair failed: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 356                         return -ECHECK
 357
 358                 return SUCCESS
 359
 360         def __find_deleteable_files (self):
 361                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 362                 DELETE_REGEX = config.get_value ('regular expressions', 'delete_regex')
 363                 dregex = re.compile (DELETE_REGEX, re.IGNORECASE)
 364
 365                 return [f for f in all_files if dregex.match (f)]
 366
 367         def __delete_list_of_files (self, dir, files, interactive=False):
 368                 # Delete a list of files
 369
 370                 assert os.path.isdir (dir)
 371
 372                 done = False
 373                 valid_y = ['Y', 'YES']
 374                 valid_n = ['N', 'NO', '']
 375
 376                 if interactive:
 377                         while not done:
 378                                 print 'Do you want to delete the following?:'
 379                                 for f in files:
 380                                         print f
 381                                 s = raw_input ('Delete [y/N]: ').upper()
 382
 383                                 if s in valid_y + valid_n:
 384                                         done = True
 385
 386                         if s in valid_n:
 387                                 return SUCCESS
 388
 389                 for f in files:
 390                         try:
 391                                 os.remove (os.path.join (dir, f))
 392                                 logger.addMessage ('Deleteing: %s' % os.path.join (dir, f), RarslaveLogger.MessageType.Debug)
 393                         except:
 394                                 logger.addMessage ('Failed to delete: %s' % os.path.join (dir, f),
 395                                                 RarslaveLogger.MessageType.Fatal)
 396                                 return -EDELETE
 397
 398                 return SUCCESS
 399
 400         def runDelete (self):
 401                 deleteable_files = self.__find_deleteable_files ()
 402                 ret = self.__delete_list_of_files (self.dir, deleteable_files, options.interactive)
 403
 404                 return ret
 405
 406         def run_all (self):
 407                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 408
 409                 # Repair Stage
 410                 ret = self.runCheckAndRepair ()
 411
 412                 if ret != SUCCESS:
 413                         logger.addMessage ('Repair stage failed for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 414                         return -ECHECK
 415
 416                 self.__update_name_matches ()
 417                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 418
 419                 # Extraction Stage
 420                 extractor = RarslaveExtractor (self.dir, self.all_p2files, \
 421                                 self.name_matched_files, self.prot_matched_files)
 422                 ret = extractor.runExtract (options.extract_dir)
 423
 424                 if ret != SUCCESS:
 425                         logger.addMessage ('Extraction stage failed for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 426                         return -EEXTRACT
 427
 428                 self.__update_name_matches ()
 429                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 430
 431                 # Deletion Stage
 432                 ret = self.runDelete ()
 433
 434                 if ret != SUCCESS:
 435                         logger.addMessage ('Deletion stage failed for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 436                         return -EDELETE
 437
 438                 logger.addMessage ('Successfully completed: %s' % self.p2file)
 439                 return SUCCESS
 440
 441 def run_command (cmd, indir=None):
 442         # Runs the specified command-line in the directory given (or, in the current directory
 443         # if none is given). It returns the status code given by the application.
 444
 445         pwd = os.getcwd ()
 446
 447         if indir != None:
 448                 assert os.path.isdir (indir) # MUST be a directory!
 449                 os.chdir (indir)
 450
 451         print 'RUNNING (%s): %s' % (indir, cmd)
 452         ret = os.system (cmd)
 453         os.chdir (pwd)
 454         return ret
 455
 456 def full_abspath (p):
 457         return os.path.abspath (os.path.expanduser (p))
 458
 459 def find_par2_files (files):
 460         """Find all par2 files in the list $files"""
 461
 462         PAR2_REGEX = config.get_value ('regular expressions', 'par2_regex')
 463         regex = re.compile (PAR2_REGEX, re.IGNORECASE)
 464         return [f for f in files if regex.match (f)]
 465
 466 def find_all_par2_files (dir):
 467         """Finds all par2 files in a directory"""
 468         # NOTE: does NOT return absolute paths
 469
 470         if not os.path.isdir (os.path.abspath (dir)):
 471                 raise ValueError # bad directory given
 472
 473         dir = os.path.abspath (dir)
 474         files = os.listdir (dir)
 475
 476         return find_par2_files (files)
 477
 478 def no_duplicates (li):
 479         """Removes all duplicates from a list"""
 480         return list(set(li))
 481
 482 def generate_all_parsets (dir):
 483         # Generate all parsets in the given directory.
 484
 485         assert os.path.isdir (dir) # Directory MUST be valid
 486
 487         parsets = []
 488         p2files = find_all_par2_files (dir)
 489
 490         for f in p2files:
 491                 p = PAR2Set (dir, f)
 492                 if p not in parsets:
 493                         parsets.append (p)
 494
 495         return parsets
 496
 497 def check_required_progs():
 498         """Check if the required programs are installed"""
 499
 500         shell_not_found = 32512
 501         needed = []
 502
 503         if run_command ('par2repair --help > /dev/null 2>&1') == shell_not_found:
 504                 needed.append ('par2repair')
 505
 506         if run_command ('unrar --help > /dev/null 2>&1') == shell_not_found:
 507                 needed.append ('unrar')
 508
 509         if run_command ('unzip --help > /dev/null 2>&1') == shell_not_found:
 510                 needed.append ('unzip')
 511
 512         if needed:
 513                 for n in needed:
 514                         print 'Needed program "%s" not found in $PATH' % (n, )
 515
 516                 sys.exit(1)
 517
 518 def run_options (options):
 519
 520         # Fix directories
 521         options.work_dir = full_abspath (options.work_dir)
 522
 523         # Make sure that the directory is valid
 524         if not os.path.isdir (options.work_dir):
 525                 sys.stderr.write ('\"%s\" is not a valid directory. Use the \"-d\"\n' % options.work_dir)
 526                 sys.stderr.write ('option to override the working directory temporarily, or edit the\n')
 527                 sys.stderr.write ('configuration file to override the working directory permanently.\n')
 528                 sys.exit (1)
 529
 530         if options.extract_dir != None:
 531                 options.extract_dir = full_abspath (options.extract_dir)
 532
 533         if options.version:
 534                 print PROGRAM + ' - ' + VERSION
 535                 print
 536                 print 'Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)'
 537                 print
 538                 print 'This program comes with ABSOLUTELY NO WARRANTY.'
 539                 print 'This is free software, and you are welcome to redistribute it'
 540                 print 'under certain conditions. See the file COPYING for details.'
 541                 sys.exit (0)
 542
 543         if options.check_progs:
 544                 check_required_progs ()
 545
 546         if options.write_def_config:
 547                 config.write_config (default=True)
 548
 549         if options.write_config:
 550                 config.write_config ()
 551
 552 def find_loglevel (options):
 553
 554         loglevel = options.verbose - options.quiet
 555
 556         if loglevel < RarslaveLogger.MessageType.Fatal:
 557                 loglevel = RarslaveLogger.MessageType.Fatal
 558
 559         if loglevel > RarslaveLogger.MessageType.Debug:
 560                 loglevel = RarslaveLogger.MessageType.Debug
 561
 562         return loglevel
 563
 564 def printMessageTable (loglevel):
 565
 566         if logger.hasFatalMessages ():
 567                 print '\nFatal Messages\n' + '=' * 80
 568                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
 569
 570         if loglevel == RarslaveLogger.MessageType.Fatal:
 571                 return
 572
 573         if logger.hasNormalMessages ():
 574                 print '\nNormal Messages\n' + '=' * 80
 575                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
 576
 577         if loglevel == RarslaveLogger.MessageType.Normal:
 578                 return
 579
 580         if logger.hasVerboseMessages ():
 581                 print '\nVerbose Messages\n' + '=' * 80
 582                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
 583
 584         if loglevel == RarslaveLogger.MessageType.Verbose:
 585                 return
 586
 587         if logger.hasDebugMessages ():
 588                 print '\nDebug Messages\n' + '=' * 80
 589                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
 590
 591         return
 592
 593 def main ():
 594
 595         # Build the OptionParser
 596         parser = optparse.OptionParser()
 597         parser.add_option('-n', '--not-recursive',
 598                                                 action='store_false', dest='recursive',
 599                                                 default=config.get_value('options', 'recursive'),
 600                                                 help="Don't run recursively")
 601
 602         parser.add_option('-d', '--work-dir',
 603                                                 dest='work_dir', type='string',
 604                                                 default=config.get_value('directories', 'working_directory'),
 605                                                 help="Start running at DIR", metavar='DIR')
 606
 607         parser.add_option('-e', '--extract-dir',
 608                                                 dest='extract_dir', type='string',
 609                                                 default=config.get_value('directories', 'extract_directory'),
 610                                                 help="Extract to DIR", metavar='DIR')
 611
 612         parser.add_option('-p', '--check-required-programs',
 613                                                 action='store_true', dest='check_progs',
 614                                                 default=False,
 615                                                 help="Check for required programs")
 616
 617         parser.add_option('-f', '--write-default-config',
 618                                                 action='store_true', dest='write_def_config',
 619                                                 default=False, help="Write out a new default config")
 620
 621         parser.add_option('-c', '--write-new-config',
 622                                                 action='store_true', dest='write_config',
 623                                                 default=False, help="Write out the current config")
 624
 625         parser.add_option('-i', '--interactive', dest='interactive', action='store_true',
 626                                                 default=config.get_value('options', 'interactive'),
 627                                                 help="Confirm before removing files")
 628
 629         parser.add_option('-q', '--quiet', dest='quiet', action='count',
 630                                                 default=0, help="Output fatal messages only")
 631
 632         parser.add_option('-v', '--verbose', dest='verbose', action='count',
 633                                                 default=0, help="Output extra information")
 634
 635         parser.add_option('-V', '--version', dest='version', action='store_true',
 636                                                 default=False, help="Output version information")
 637
 638         parser.version = VERSION
 639
 640         # Parse the given options
 641         global options
 642         (options, args) = parser.parse_args()
 643
 644         # Run any special actions that are needed on these options
 645         run_options (options)
 646
 647         # Find the loglevel using the options given
 648         loglevel = find_loglevel (options)
 649
 650         # Run recursively
 651         if options.recursive:
 652                 for (dir, subdirs, files) in os.walk (options.work_dir):
 653                         parsets = generate_all_parsets (dir)
 654                         for p in parsets:
 655                                 p.run_all ()
 656
 657         # Non-recursive
 658         else:
 659                 parsets = generate_all_parsets (options.work_dir)
 660                 for p in parsets:
 661                         p.run_all ()
 662
 663         # Print the results
 664         printMessageTable (loglevel)
 665
 666         # Done!
 667         return 0
 668
 669 if __name__ == '__main__':
 670         main ()
 671