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