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 __is_joinfile (self, filename):
 325                 regex = re.compile ('^.*\.\d\d\d$', re.IGNORECASE)
 326                 if regex.match (filename):
 327                         return True
 328
 329                 return False
 330
 331         def __should_be_joined (self, files):
 332                 for f in files:
 333                         if self.__is_joinfile (f):
 334                                 return True
 335
 336         def runCheckAndRepair (self):
 337                 PAR2_CMD = config.get_value ('commands', 'par2repair')
 338
 339                 # Get set up
 340                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 341                 join = self.__should_be_joined (all_files)
 342
 343                 # assemble the command
 344                 # par2repair -- PAR2 PAR2_EXTRA [JOIN_FILES]
 345                 command = "%s \"%s\" " % (PAR2_CMD, self.p2file)
 346
 347                 for f in self.all_p2files:
 348                         if f != self.p2file:
 349                                 command += "\"%s\" " % os.path.split (f)[1]
 350
 351                 # Only needed when using par2 to join
 352                 if join:
 353                         for f in all_files:
 354                                 if self.__is_joinfile (f):
 355                                         command += "\"%s\" " % os.path.split (f)[1]
 356
 357                 # run the command
 358                 ret = run_command (command, self.dir)
 359
 360                 # check the result
 361                 if ret != 0:
 362                         logger.addMessage ('PAR2 Check / Repair failed: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 363                         return -ECHECK
 364
 365                 return SUCCESS
 366
 367         def __find_deleteable_files (self):
 368                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 369                 DELETE_REGEX = config.get_value ('regular expressions', 'delete_regex')
 370                 dregex = re.compile (DELETE_REGEX, re.IGNORECASE)
 371
 372                 return [f for f in all_files if dregex.match (f)]
 373
 374         def __delete_list_of_files (self, dir, files, interactive=False):
 375                 # Delete a list of files
 376
 377                 assert os.path.isdir (dir)
 378
 379                 done = False
 380                 valid_y = ['Y', 'YES']
 381                 valid_n = ['N', 'NO', '']
 382
 383                 if interactive:
 384                         while not done:
 385                                 print 'Do you want to delete the following?:'
 386                                 for f in files:
 387                                         print f
 388                                 s = raw_input ('Delete [y/N]: ').upper()
 389
 390                                 if s in valid_y + valid_n:
 391                                         done = True
 392
 393                         if s in valid_n:
 394                                 return SUCCESS
 395
 396                 for f in files:
 397                         try:
 398                                 os.remove (os.path.join (dir, f))
 399                                 logger.addMessage ('Deleteing: %s' % os.path.join (dir, f), RarslaveLogger.MessageType.Debug)
 400                         except:
 401                                 logger.addMessage ('Failed to delete: %s' % os.path.join (dir, f),
 402                                                 RarslaveLogger.MessageType.Fatal)
 403                                 return -EDELETE
 404
 405                 return SUCCESS
 406
 407         def runDelete (self):
 408                 deleteable_files = self.__find_deleteable_files ()
 409                 ret = self.__delete_list_of_files (self.dir, deleteable_files, options.interactive)
 410
 411                 return ret
 412
 413         def run_all (self):
 414                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 415
 416                 # Repair Stage
 417                 ret = self.runCheckAndRepair ()
 418
 419                 if ret != SUCCESS:
 420                         logger.addMessage ('Repair stage failed for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 421                         return -ECHECK
 422
 423                 self.__update_name_matches ()
 424                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 425
 426                 # Extraction Stage
 427                 extractor = RarslaveExtractor (self.dir, self.all_p2files, \
 428                                 self.name_matched_files, self.prot_matched_files)
 429                 ret = extractor.runExtract (options.extract_dir)
 430
 431                 if ret != SUCCESS:
 432                         logger.addMessage ('Extraction stage failed for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 433                         return -EEXTRACT
 434
 435                 self.__update_name_matches ()
 436                 all_files = no_duplicates (self.name_matched_files + self.prot_matched_files)
 437
 438                 # Deletion Stage
 439                 ret = self.runDelete ()
 440
 441                 if ret != SUCCESS:
 442                         logger.addMessage ('Deletion stage failed for: %s' % self.p2file, RarslaveLogger.MessageType.Fatal)
 443                         return -EDELETE
 444
 445                 logger.addMessage ('Successfully completed: %s' % self.p2file)
 446                 return SUCCESS
 447
 448 def run_command (cmd, indir=None):
 449         # Runs the specified command-line in the directory given (or, in the current directory
 450         # if none is given). It returns the status code given by the application.
 451
 452         pwd = os.getcwd ()
 453
 454         if indir != None:
 455                 assert os.path.isdir (indir) # MUST be a directory!
 456                 os.chdir (indir)
 457
 458         ret = os.system (cmd)
 459         os.chdir (pwd)
 460         return ret
 461
 462 def full_abspath (p):
 463         return os.path.abspath (os.path.expanduser (p))
 464
 465 def find_par2_files (files):
 466         """Find all par2 files in the list $files"""
 467
 468         PAR2_REGEX = config.get_value ('regular expressions', 'par2_regex')
 469         regex = re.compile (PAR2_REGEX, re.IGNORECASE)
 470         return [f for f in files if regex.match (f)]
 471
 472 def find_all_par2_files (dir):
 473         """Finds all par2 files in a directory"""
 474         # NOTE: does NOT return absolute paths
 475
 476         if not os.path.isdir (os.path.abspath (dir)):
 477                 raise ValueError # bad directory given
 478
 479         dir = os.path.abspath (dir)
 480         files = os.listdir (dir)
 481
 482         return find_par2_files (files)
 483
 484 def no_duplicates (li):
 485         """Removes all duplicates from a list"""
 486         return list(set(li))
 487
 488 def generate_all_parsets (dir):
 489         # Generate all parsets in the given directory.
 490
 491         assert os.path.isdir (dir) # Directory MUST be valid
 492
 493         parsets = []
 494         p2files = find_all_par2_files (dir)
 495
 496         for f in p2files:
 497                 p = PAR2Set (dir, f)
 498                 if p not in parsets:
 499                         parsets.append (p)
 500
 501         return parsets
 502
 503 def check_required_progs():
 504         """Check if the required programs are installed"""
 505
 506         shell_not_found = 32512
 507         needed = []
 508
 509         if run_command ('par2repair --help > /dev/null 2>&1') == shell_not_found:
 510                 needed.append ('par2repair')
 511
 512         if run_command ('unrar --help > /dev/null 2>&1') == shell_not_found:
 513                 needed.append ('unrar')
 514
 515         if run_command ('unzip --help > /dev/null 2>&1') == shell_not_found:
 516                 needed.append ('unzip')
 517
 518         if needed:
 519                 for n in needed:
 520                         print 'Needed program "%s" not found in $PATH' % (n, )
 521
 522                 sys.exit(1)
 523
 524 def run_options (options):
 525
 526         # Fix directories
 527         options.work_dir = full_abspath (options.work_dir)
 528
 529         # Make sure that the directory is valid
 530         if not os.path.isdir (options.work_dir):
 531                 sys.stderr.write ('\"%s\" is not a valid directory. Use the \"-d\"\n' % options.work_dir)
 532                 sys.stderr.write ('option to override the working directory temporarily, or edit the\n')
 533                 sys.stderr.write ('configuration file to override the working directory permanently.\n')
 534                 sys.exit (1)
 535
 536         if options.extract_dir != None:
 537                 options.extract_dir = full_abspath (options.extract_dir)
 538
 539         if options.version:
 540                 print PROGRAM + ' - ' + VERSION
 541                 print
 542                 print 'Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)'
 543                 print
 544                 print 'This program comes with ABSOLUTELY NO WARRANTY.'
 545                 print 'This is free software, and you are welcome to redistribute it'
 546                 print 'under certain conditions. See the file COPYING for details.'
 547                 sys.exit (0)
 548
 549         if options.check_progs:
 550                 check_required_progs ()
 551
 552         if options.write_def_config:
 553                 config.write_config (default=True)
 554
 555         if options.write_config:
 556                 config.write_config ()
 557
 558 def find_loglevel (options):
 559
 560         loglevel = options.verbose - options.quiet
 561
 562         if loglevel < RarslaveLogger.MessageType.Fatal:
 563                 loglevel = RarslaveLogger.MessageType.Fatal
 564
 565         if loglevel > RarslaveLogger.MessageType.Debug:
 566                 loglevel = RarslaveLogger.MessageType.Debug
 567
 568         return loglevel
 569
 570 def printMessageTable (loglevel):
 571
 572         if logger.hasFatalMessages ():
 573                 print '\nFatal Messages\n' + '=' * 80
 574                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
 575
 576         if loglevel == RarslaveLogger.MessageType.Fatal:
 577                 return
 578
 579         if logger.hasNormalMessages ():
 580                 print '\nNormal Messages\n' + '=' * 80
 581                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
 582
 583         if loglevel == RarslaveLogger.MessageType.Normal:
 584                 return
 585
 586         if logger.hasVerboseMessages ():
 587                 print '\nVerbose Messages\n' + '=' * 80
 588                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
 589
 590         if loglevel == RarslaveLogger.MessageType.Verbose:
 591                 return
 592
 593         if logger.hasDebugMessages ():
 594                 print '\nDebug Messages\n' + '=' * 80
 595                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
 596
 597         return
 598
 599 def main ():
 600
 601         # Build the OptionParser
 602         parser = optparse.OptionParser()
 603         parser.add_option('-n', '--not-recursive',
 604                                                 action='store_false', dest='recursive',
 605                                                 default=config.get_value('options', 'recursive'),
 606                                                 help="Don't run recursively")
 607
 608         parser.add_option('-d', '--work-dir',
 609                                                 dest='work_dir', type='string',
 610                                                 default=config.get_value('directories', 'working_directory'),
 611                                                 help="Start running at DIR", metavar='DIR')
 612
 613         parser.add_option('-e', '--extract-dir',
 614                                                 dest='extract_dir', type='string',
 615                                                 default=config.get_value('directories', 'extract_directory'),
 616                                                 help="Extract to DIR", metavar='DIR')
 617
 618         parser.add_option('-p', '--check-required-programs',
 619                                                 action='store_true', dest='check_progs',
 620                                                 default=False,
 621                                                 help="Check for required programs")
 622
 623         parser.add_option('-f', '--write-default-config',
 624                                                 action='store_true', dest='write_def_config',
 625                                                 default=False, help="Write out a new default config")
 626
 627         parser.add_option('-c', '--write-new-config',
 628                                                 action='store_true', dest='write_config',
 629                                                 default=False, help="Write out the current config")
 630
 631         parser.add_option('-i', '--interactive', dest='interactive', action='store_true',
 632                                                 default=config.get_value('options', 'interactive'),
 633                                                 help="Confirm before removing files")
 634
 635         parser.add_option('-q', '--quiet', dest='quiet', action='count',
 636                                                 default=0, help="Output fatal messages only")
 637
 638         parser.add_option('-v', '--verbose', dest='verbose', action='count',
 639                                                 default=0, help="Output extra information")
 640
 641         parser.add_option('-V', '--version', dest='version', action='store_true',
 642                                                 default=False, help="Output version information")
 643
 644         parser.version = VERSION
 645
 646         # Parse the given options
 647         global options
 648         (options, args) = parser.parse_args()
 649
 650         # Run any special actions that are needed on these options
 651         run_options (options)
 652
 653         # Find the loglevel using the options given
 654         loglevel = find_loglevel (options)
 655
 656         # Run recursively
 657         if options.recursive:
 658                 for (dir, subdirs, files) in os.walk (options.work_dir):
 659                         parsets = generate_all_parsets (dir)
 660                         for p in parsets:
 661                                 p.run_all ()
 662
 663         # Non-recursive
 664         else:
 665                 parsets = generate_all_parsets (options.work_dir)
 666                 for p in parsets:
 667                         p.run_all ()
 668
 669         # Print the results
 670         printMessageTable (loglevel)
 671
 672         # Done!
 673         return 0
 674
 675 if __name__ == '__main__':
 676         main ()
 677