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