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         options.work_dir = full_abspath (options.work_dir)
 483
 484         if options.extract_dir != None:
 485                 options.extract_dir = full_abspath (options.extract_dir)
 486
 487         if options.version:
 488                 print PROGRAM + ' - ' + VERSION
 489                 print
 490                 print 'Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)'
 491                 print
 492                 print 'This program comes with ABSOLUTELY NO WARRANTY.'
 493                 print 'This is free software, and you are welcome to redistribute it'
 494                 print 'under certain conditions. See the file COPYING for details.'
 495                 sys.exit (0)
 496
 497         if options.check_progs:
 498                 check_required_progs ()
 499
 500         if options.write_def_config:
 501                 config.write_config (default=True)
 502
 503         if options.write_config:
 504                 config.write_config ()
 505
 506 def find_loglevel (options):
 507
 508         loglevel = options.verbose - options.quiet
 509
 510         if loglevel < RarslaveLogger.MessageType.Fatal:
 511                 loglevel = RarslaveLogger.MessageType.Fatal
 512
 513         if loglevel > RarslaveLogger.MessageType.Debug:
 514                 loglevel = RarslaveLogger.MessageType.Debug
 515
 516         return loglevel
 517
 518 def printMessageTable (loglevel):
 519
 520         if logger.hasFatalMessages ():
 521                 print '\nFatal Messages\n' + '=' * 80
 522                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
 523
 524         if loglevel == RarslaveLogger.MessageType.Fatal:
 525                 return
 526
 527         if logger.hasNormalMessages ():
 528                 print '\nNormal Messages\n' + '=' * 80
 529                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
 530
 531         if loglevel == RarslaveLogger.MessageType.Normal:
 532                 return
 533
 534         if logger.hasVerboseMessages ():
 535                 print '\nVerbose Messages\n' + '=' * 80
 536                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
 537
 538         if loglevel == RarslaveLogger.MessageType.Verbose:
 539                 return
 540
 541         if logger.hasDebugMessages ():
 542                 print '\nDebug Messages\n' + '=' * 80
 543                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
 544
 545         return
 546
 547 def main ():
 548
 549         # Build the OptionParser
 550         parser = optparse.OptionParser()
 551         parser.add_option('-n', '--not-recursive',
 552                                                 action='store_false', dest='recursive',
 553                                                 default=config.get_value('options', 'recursive'),
 554                                                 help="Don't run recursively")
 555
 556         parser.add_option('-d', '--work-dir',
 557                                                 dest='work_dir', type='string',
 558                                                 default=config.get_value('directories', 'working_directory'),
 559                                                 help="Start running at DIR", metavar='DIR')
 560
 561         parser.add_option('-e', '--extract-dir',
 562                                                 dest='extract_dir', type='string',
 563                                                 default=config.get_value('directories', 'extract_directory'),
 564                                                 help="Extract to DIR", metavar='DIR')
 565
 566         parser.add_option('-p', '--check-required-programs',
 567                                                 action='store_true', dest='check_progs',
 568                                                 default=False,
 569                                                 help="Check for required programs")
 570
 571         parser.add_option('-f', '--write-default-config',
 572                                                 action='store_true', dest='write_def_config',
 573                                                 default=False, help="Write out a new default config")
 574
 575         parser.add_option('-c', '--write-new-config',
 576                                                 action='store_true', dest='write_config',
 577                                                 default=False, help="Write out the current config")
 578
 579         parser.add_option('-i', '--interactive', dest='interactive', action='store_true',
 580                                                 default=config.get_value('options', 'interactive'),
 581                                                 help="Confirm before removing files")
 582
 583         parser.add_option('-q', '--quiet', dest='quiet', action='count',
 584                                                 default=0, help="Output fatal messages only")
 585
 586         parser.add_option('-v', '--verbose', dest='verbose', action='count',
 587                                                 default=0, help="Output extra information")
 588
 589         parser.add_option('-V', '--version', dest='version', action='store_true',
 590                                                 default=False, help="Output version information")
 591
 592         parser.version = VERSION
 593
 594         # Parse the given options
 595         global options
 596         (options, args) = parser.parse_args()
 597
 598         # Run any special actions that are needed on these options
 599         run_options (options)
 600
 601         # Find the loglevel using the options given
 602         loglevel = find_loglevel (options)
 603
 604         # Run recursively
 605         if options.recursive:
 606                 for (dir, subdirs, files) in os.walk (options.work_dir):
 607                         parsets = generate_all_parsets (dir)
 608                         for p in parsets:
 609                                 p.run_all ()
 610
 611         # Non-recursive
 612         else:
 613                 parsets = generate_all_parsets (options.work_dir)
 614                 for p in parsets:
 615                         p.run_all ()
 616
 617         # Print the results
 618         printMessageTable (loglevel)
 619
 620         # Done!
 621         return 0
 622
 623 if __name__ == '__main__':
 624         main ()
 625