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