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         print 'RUNNING (%s): %s' % (indir, cmd)
 459         ret = os.system (cmd)
 460         os.chdir (pwd)
 461         return ret
 462
 463 def full_abspath (p):
 464         return os.path.abspath (os.path.expanduser (p))
 465
 466 def find_par2_files (files):
 467         """Find all par2 files in the list $files"""
 468
 469         PAR2_REGEX = config.get_value ('regular expressions', 'par2_regex')
 470         regex = re.compile (PAR2_REGEX, re.IGNORECASE)
 471         return [f for f in files if regex.match (f)]
 472
 473 def find_all_par2_files (dir):
 474         """Finds all par2 files in a directory"""
 475         # NOTE: does NOT return absolute paths
 476
 477         if not os.path.isdir (os.path.abspath (dir)):
 478                 raise ValueError # bad directory given
 479
 480         dir = os.path.abspath (dir)
 481         files = os.listdir (dir)
 482
 483         return find_par2_files (files)
 484
 485 def no_duplicates (li):
 486         """Removes all duplicates from a list"""
 487         return list(set(li))
 488
 489 def generate_all_parsets (dir):
 490         # Generate all parsets in the given directory.
 491
 492         assert os.path.isdir (dir) # Directory MUST be valid
 493
 494         parsets = []
 495         p2files = find_all_par2_files (dir)
 496
 497         for f in p2files:
 498                 p = PAR2Set (dir, f)
 499                 if p not in parsets:
 500                         parsets.append (p)
 501
 502         return parsets
 503
 504 def check_required_progs():
 505         """Check if the required programs are installed"""
 506
 507         shell_not_found = 32512
 508         needed = []
 509
 510         if run_command ('par2repair --help > /dev/null 2>&1') == shell_not_found:
 511                 needed.append ('par2repair')
 512
 513         if run_command ('unrar --help > /dev/null 2>&1') == shell_not_found:
 514                 needed.append ('unrar')
 515
 516         if run_command ('unzip --help > /dev/null 2>&1') == shell_not_found:
 517                 needed.append ('unzip')
 518
 519         if needed:
 520                 for n in needed:
 521                         print 'Needed program "%s" not found in $PATH' % (n, )
 522
 523                 sys.exit(1)
 524
 525 def run_options (options):
 526
 527         # Fix directories
 528         options.work_dir = full_abspath (options.work_dir)
 529
 530         # Make sure that the directory is valid
 531         if not os.path.isdir (options.work_dir):
 532                 sys.stderr.write ('\"%s\" is not a valid directory. Use the \"-d\"\n' % options.work_dir)
 533                 sys.stderr.write ('option to override the working directory temporarily, or edit the\n')
 534                 sys.stderr.write ('configuration file to override the working directory permanently.\n')
 535                 sys.exit (1)
 536
 537         if options.extract_dir != None:
 538                 options.extract_dir = full_abspath (options.extract_dir)
 539
 540         if options.version:
 541                 print PROGRAM + ' - ' + VERSION
 542                 print
 543                 print 'Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)'
 544                 print
 545                 print 'This program comes with ABSOLUTELY NO WARRANTY.'
 546                 print 'This is free software, and you are welcome to redistribute it'
 547                 print 'under certain conditions. See the file COPYING for details.'
 548                 sys.exit (0)
 549
 550         if options.check_progs:
 551                 check_required_progs ()
 552
 553         if options.write_def_config:
 554                 config.write_config (default=True)
 555
 556         if options.write_config:
 557                 config.write_config ()
 558
 559 def find_loglevel (options):
 560
 561         loglevel = options.verbose - options.quiet
 562
 563         if loglevel < RarslaveLogger.MessageType.Fatal:
 564                 loglevel = RarslaveLogger.MessageType.Fatal
 565
 566         if loglevel > RarslaveLogger.MessageType.Debug:
 567                 loglevel = RarslaveLogger.MessageType.Debug
 568
 569         return loglevel
 570
 571 def printMessageTable (loglevel):
 572
 573         if logger.hasFatalMessages ():
 574                 print '\nFatal Messages\n' + '=' * 80
 575                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
 576
 577         if loglevel == RarslaveLogger.MessageType.Fatal:
 578                 return
 579
 580         if logger.hasNormalMessages ():
 581                 print '\nNormal Messages\n' + '=' * 80
 582                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
 583
 584         if loglevel == RarslaveLogger.MessageType.Normal:
 585                 return
 586
 587         if logger.hasVerboseMessages ():
 588                 print '\nVerbose Messages\n' + '=' * 80
 589                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
 590
 591         if loglevel == RarslaveLogger.MessageType.Verbose:
 592                 return
 593
 594         if logger.hasDebugMessages ():
 595                 print '\nDebug Messages\n' + '=' * 80
 596                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
 597
 598         return
 599
 600 def main ():
 601
 602         # Build the OptionParser
 603         parser = optparse.OptionParser()
 604         parser.add_option('-n', '--not-recursive',
 605                                                 action='store_false', dest='recursive',
 606                                                 default=config.get_value('options', 'recursive'),
 607                                                 help="Don't run recursively")
 608
 609         parser.add_option('-d', '--work-dir',
 610                                                 dest='work_dir', type='string',
 611                                                 default=config.get_value('directories', 'working_directory'),
 612                                                 help="Start running at DIR", metavar='DIR')
 613
 614         parser.add_option('-e', '--extract-dir',
 615                                                 dest='extract_dir', type='string',
 616                                                 default=config.get_value('directories', 'extract_directory'),
 617                                                 help="Extract to DIR", metavar='DIR')
 618
 619         parser.add_option('-p', '--check-required-programs',
 620                                                 action='store_true', dest='check_progs',
 621                                                 default=False,
 622                                                 help="Check for required programs")
 623
 624         parser.add_option('-f', '--write-default-config',
 625                                                 action='store_true', dest='write_def_config',
 626                                                 default=False, help="Write out a new default config")
 627
 628         parser.add_option('-c', '--write-new-config',
 629                                                 action='store_true', dest='write_config',
 630                                                 default=False, help="Write out the current config")
 631
 632         parser.add_option('-i', '--interactive', dest='interactive', action='store_true',
 633                                                 default=config.get_value('options', 'interactive'),
 634                                                 help="Confirm before removing files")
 635
 636         parser.add_option('-q', '--quiet', dest='quiet', action='count',
 637                                                 default=0, help="Output fatal messages only")
 638
 639         parser.add_option('-v', '--verbose', dest='verbose', action='count',
 640                                                 default=0, help="Output extra information")
 641
 642         parser.add_option('-V', '--version', dest='version', action='store_true',
 643                                                 default=False, help="Output version information")
 644
 645         parser.version = VERSION
 646
 647         # Parse the given options
 648         global options
 649         (options, args) = parser.parse_args()
 650
 651         # Run any special actions that are needed on these options
 652         run_options (options)
 653
 654         # Find the loglevel using the options given
 655         loglevel = find_loglevel (options)
 656
 657         # Run recursively
 658         if options.recursive:
 659                 for (dir, subdirs, files) in os.walk (options.work_dir):
 660                         parsets = generate_all_parsets (dir)
 661                         for p in parsets:
 662                                 p.run_all ()
 663
 664         # Non-recursive
 665         else:
 666                 parsets = generate_all_parsets (options.work_dir)
 667                 for p in parsets:
 668                         p.run_all ()
 669
 670         # Print the results
 671         printMessageTable (loglevel)
 672
 673         # Done!
 674         return 0
 675
 676 if __name__ == '__main__':
 677         main ()
 678