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