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