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