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