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