rarslave.py

   1 #!/usr/bin/env python
   2 # vim: set ts=4 sts=4 sw=4 textwidth=112 :
   3
   4 import re, os, sys
   5 import par2parser
   6 import RarslaveConfig
   7 from RarslaveLogger import RarslaveLogger
   8
   9 # Global Variables
  10 (TYPE_OLDRAR, TYPE_NEWRAR, TYPE_ZIP, TYPE_NOEXTRACT) = range (4)
  11 (SUCCESS, ECHECK, EEXTRACT, EDELETE) = range(4)
  12 config = RarslaveConfig.RarslaveConfig()
  13 logger = RarslaveLogger ()
  14
  15 class RarslaveExtractor (object):
  16
  17         def __init__ (self, type):
  18                 self.type = type
  19                 self.heads = []
  20
  21         def addHead (self, dir, head):
  22                 assert os.path.isdir (dir)
  23                 # REQUIRES that the dir is valid, but not that the file is valid, so that
  24                 # we can move a file that doesn't exist yet.
  25                 # FIXME: probably CAN add this back, since we should be running this AFTER repair.
  26                 #assert os.path.isfile (os.path.join (dir, head))
  27
  28                 self.heads.append (os.path.join (dir, head))
  29
  30         def extract (self, todir=None):
  31                 # Extract all heads of this set
  32
  33                 # Create the directory $todir if it doesn't exist
  34                 if todir != None and not os.path.isdir (todir):
  35                         logger.addMessage ('Creating directory: %s' % todir, False)
  36                         try:
  37                                 os.makedirs (todir)
  38                         except OSError:
  39                                 logger.addMessage ('FAILED to create directory: %s' % todir)
  40                                 return -EEXTRACT
  41
  42                 # Extract all heads
  43                 extraction_func = \
  44                         { TYPE_OLDRAR : self.__extract_rar,
  45                           TYPE_NEWRAR : self.__extract_rar,
  46                           TYPE_ZIP    : self.__extract_zip,
  47                           TYPE_NOEXTRACT : self.__extract_noextract }[self.type]
  48
  49                 # Call the extraction function on each head
  50                 for h in self.heads:
  51                         if todir == None:
  52                                 # Run in the head's directory
  53                                 ret = extraction_func (h, os.path.dirname (h))
  54                         else:
  55                                 ret = extraction_func (h, todir)
  56
  57                         # Check error code
  58                         if ret != SUCCESS:
  59                                 logger.addMessage ('Failed extracting: %s' % h)
  60                                 return -EEXTRACT
  61
  62                 return SUCCESS
  63
  64         def __extract_rar (self, file, todir):
  65                 assert os.path.isfile (file)
  66                 assert os.path.isdir (todir)
  67
  68                 RAR_CMD = config.get_value ('commands', 'unrar')
  69
  70                 cmd = '%s \"%s\"' % (RAR_CMD, file)
  71                 ret = run_command (cmd, todir)
  72
  73                 # Check error code
  74                 if ret != 0:
  75                         return -EEXTRACT
  76
  77                 return SUCCESS
  78
  79         def __extract_zip (self, file, todir):
  80                 ZIP_CMD = config.get_value ('commands', 'unzip')
  81
  82                 cmd = ZIP_CMD % (file, todir)
  83                 ret = run_command (cmd)
  84
  85                 # Check error code
  86                 if ret != 0:
  87                         return -EEXTRACT
  88
  89                 return SUCCESS
  90
  91         def __extract_noextract (self, file, todir):
  92                 # Just move this file to the $todir, since no extraction is needed
  93                 # FIXME: NOTE: mv will fail by itself if you're moving to the same dir!
  94                 NOEXTRACT_CMD = config.get_value ('commands', 'noextract')
  95
  96                 cmd = NOEXTRACT_CMD % (file, todir)
  97                 ret = run_command (cmd)
  98
  99                 # Check error code
 100                 if ret != 0:
 101                         return -EEXTRACT
 102
 103                 return SUCCESS
 104
 105
 106
 107 class RarslaveRepairer (object):
 108         # Verify (and repair) the set
 109         # Make sure it worked, otherwise clean up and return failure
 110
 111         def __init__ (self, dir, file, join=False):
 112                 self.dir  = dir  # the directory containing the par2 file
 113                 self.file = file # the par2 file
 114                 self.join = join # True if the par2 set is 001 002 ...
 115
 116                 assert os.path.isdir (dir)
 117                 assert os.path.isfile (os.path.join (dir, file))
 118
 119         def checkAndRepair (self):
 120                 # Form the command:
 121                 # par2repair -- PAR2 PAR2_EXTRA [JOIN_FILES]
 122                 PAR2_CMD = config.get_value ('commands', 'par2repair')
 123
 124                 # Get set up
 125                 basename = get_basename (self.file)
 126                 all_files = find_likely_files (basename, self.dir)
 127                 all_files.sort ()
 128                 par2_files = find_par2_files (all_files)
 129
 130                 # assemble the command
 131                 command = "%s \"%s\" " % (PAR2_CMD, self.file)
 132
 133                 for f in par2_files:
 134                         if f != self.file:
 135                                 command += "\"%s\" " % os.path.split (f)[1]
 136
 137                 if self.join:
 138                         for f in all_files:
 139                                 if f not in par2_files:
 140                                         command += "\"%s\" " % os.path.split (f)[1]
 141
 142                 # run the command
 143                 ret = run_command (command, self.dir)
 144
 145                 # check the result
 146                 if ret != 0:
 147                         logger.addMessage ('PAR2 Check / Repair failed: %s' % self.file)
 148                         return -ECHECK
 149
 150                 return SUCCESS
 151
 152 def run_command (cmd, indir=None):
 153         # Runs the specified command-line in the directory given (or, in the current directory
 154         # if none is given). It returns the status code given by the application.
 155
 156         pwd = os.getcwd ()
 157
 158         if indir != None:
 159                 assert os.path.isdir (indir) # MUST be a directory!
 160                 os.chdir (indir)
 161
 162         # FIXME: re-enable this after testing
 163         print 'RUNNING (%s): %s' % (indir, cmd)
 164         return SUCCESS
 165
 166         # return os.system (cmd)
 167
 168
 169 def full_abspath (p):
 170         return os.path.abspath (os.path.expanduser (p))
 171
 172 def get_basename (name):
 173         """Strips most kinds of endings from a filename"""
 174
 175         regex = config.get_value ('regular expressions', 'basename_regex')
 176         r = re.compile (regex, re.IGNORECASE)
 177         done = False
 178
 179         while not done:
 180                 done = True
 181
 182                 if r.match (name):
 183                         g = r.match (name).groups()
 184                         name = g[0]
 185                         done = False
 186
 187         return name
 188
 189 def find_likely_files (name, dir):
 190         """Finds files which are likely to be part of the set corresponding
 191            to $name in the directory $dir"""
 192
 193         if not os.path.isdir (os.path.abspath (dir)):
 194                 raise ValueError # bad directory given
 195
 196         dir = os.path.abspath (dir)
 197         ename = re.escape (name)
 198         regex = re.compile ('^%s.*$' % (ename, ))
 199
 200         return [f for f in os.listdir (dir) if regex.match (f)]
 201
 202 def find_par2_files (files):
 203         """Find all par2 files in the list $files"""
 204
 205         PAR2_REGEX = config.get_value ('regular expressions', 'par2_regex')
 206         regex = re.compile (PAR2_REGEX, re.IGNORECASE)
 207         return [f for f in files if regex.match (f)]
 208
 209 def find_all_par2_files (dir):
 210         """Finds all par2 files in a directory"""
 211         # NOTE: does NOT return absolute paths
 212
 213         if not os.path.isdir (os.path.abspath (dir)):
 214                 raise ValueError # bad directory given
 215
 216         dir = os.path.abspath (dir)
 217         files = os.listdir (dir)
 218
 219         return find_par2_files (files)
 220
 221 def has_extension (f, ext):
 222         """Checks if f has the extension ext"""
 223
 224         if ext[0] != '.':
 225                 ext = '.' + ext
 226
 227         ext = re.escape (ext)
 228         regex = re.compile ('^.*%s$' % (ext, ), re.IGNORECASE)
 229         return regex.match (f)
 230
 231 def find_extraction_heads (dir, files):
 232         """Takes a list of possible files and finds likely heads of
 233            extraction."""
 234
 235         # NOTE: perhaps this should happen AFTER repair is
 236         # NOTE: successful. That way all files would already exist
 237
 238         # According to various sources online:
 239         # 1) pre rar-3.0: .rar .r00 .r01 ...
 240         # 2) post rar-3.0: .part01.rar .part02.rar
 241         # 3) zip all ver: .zip
 242
 243         extractor = None
 244         p2files = find_par2_files (files)
 245
 246         # Old RAR type, find all files ending in .rar
 247         if is_oldrar (files):
 248                 extractor = RarslaveExtractor (TYPE_OLDRAR)
 249                 regex = re.compile ('^.*\.rar$', re.IGNORECASE)
 250                 for f in files:
 251                         if regex.match (f):
 252                                 extractor.addHead (dir, f)
 253
 254         if is_newrar (files):
 255                 extractor = RarslaveExtractor (TYPE_NEWRAR)
 256                 regex = re.compile ('^.*\.part01.rar$', re.IGNORECASE)
 257                 for f in files:
 258                         if regex.match (f):
 259                                 extractor.addHead (dir, f)
 260
 261         if is_zip (files):
 262                 extractor = RarslaveExtractor (TYPE_ZIP)
 263                 regex = re.compile ('^.*\.zip$', re.IGNORECASE)
 264                 for f in files:
 265                         if regex.match (f):
 266                                 extractor.addHead (dir, f)
 267
 268         if is_noextract (files):
 269                 # Use the Par2 Parser (from cfv) here to find out what files are protected.
 270                 # Since these are not being extracted, they will be mv'd to another directory
 271                 # later.
 272                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
 273
 274                 for f in p2files:
 275                         done = False
 276                         try:
 277                                 prot_files = par2parser.get_protected_files (dir, f)
 278                                 done = True
 279                         except: #FIXME: add the actual exceptions
 280                                 logger.addMessage ('Error parsing PAR2 file: %s', f)
 281                                 continue
 282
 283                         if done:
 284                                 break
 285
 286                 if done:
 287                         for f in prot_files:
 288                                 extractor.addHead (dir, f)
 289                 else:
 290                         logger.addMessage ('Error parsing all PAR2 files in this set ...', True)
 291
 292         # Make sure we found the type
 293         if extractor == None:
 294                 logger.addMessage ('Not able to find an extractor for this type of set: %s' % p2files[0])
 295
 296                 # No-heads here, but it's better than failing completely
 297                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
 298
 299         return extractor
 300
 301 def is_oldrar (files):
 302         for f in files:
 303                 if has_extension (f, '.r00'):
 304                         return True
 305
 306         return False
 307
 308 def is_newrar (files):
 309         for f in files:
 310                 if has_extension (f, '.part01.rar'):
 311                         return True
 312
 313         return False
 314
 315 def is_zip (files):
 316         for f in files:
 317                 if has_extension (f, '.zip'):
 318                         return True
 319
 320         return False
 321
 322 def is_noextract (files):
 323         # Type that needs no extraction.
 324         # TODO: Add others ???
 325         for f in files:
 326                 if has_extension (f, '.001'):
 327                         return True
 328
 329         return False
 330
 331 def find_deleteable_files (files):
 332         # Deleteable types regex should come from the config
 333         dfiles = []
 334         DELETE_REGEX = config.get_value ('regular expressions', 'delete_regex')
 335         dregex = re.compile (DELETE_REGEX, re.IGNORECASE)
 336
 337         return [f for f in files if dregex.match (f)]
 338
 339 def printlist (li):
 340         for f in li:
 341                 print f
 342
 343 class PAR2Set (object):
 344
 345         dir = None
 346         file = None
 347         likely_files = []
 348
 349         def __init__ (self, dir, file):
 350                 assert os.path.isdir (dir)
 351                 assert os.path.isfile (os.path.join (dir, file))
 352
 353                 self.dir = dir
 354                 self.file = file
 355
 356                 basename = get_basename (file)
 357                 self.likely_files = find_likely_files (basename, dir)
 358
 359         def __list_eq (self, l1, l2):
 360
 361                 if len(l1) != len(l2):
 362                         return False
 363
 364                 for e in l1:
 365                         if e not in l2:
 366                                 return False
 367
 368                 return True
 369
 370         def __eq__ (self, rhs):
 371                 return self.__list_eq (self.likely_files, rhs.likely_files)
 372
 373         def run_all (self):
 374                 par2files = find_par2_files (self.likely_files)
 375                 par2head = par2files[0]
 376
 377                 join = is_noextract (self.likely_files)
 378
 379                 # Repair Stage
 380                 repairer = RarslaveRepairer (self.dir, par2head, join)
 381                 ret = repairer.checkAndRepair ()
 382
 383                 if ret != SUCCESS:
 384                         logger.addMessage ('Repair stage failed for: %s' % par2head)
 385                         return -ECHECK
 386
 387                 # Extraction Stage
 388                 EXTRACT_DIR = config.get_value ('directories', 'extract_directory')
 389                 extractor = find_extraction_heads (self.dir, self.likely_files)
 390                 ret = extractor.extract (EXTRACT_DIR)
 391
 392                 if ret != SUCCESS:
 393                         logger.addMessage ('Extraction stage failed for: %s' % par2head)
 394                         return -EEXTRACT
 395
 396                 # Deletion Stage
 397                 DELETE_INTERACTIVE = config.get_value ('options', 'interactive')
 398                 deleteable_files = find_deleteable_files (self.likely_files)
 399                 ret = delete_list (deleteable_files, DELETE_INTERACTIVE)
 400
 401                 if ret != SUCCESS:
 402                         logger.addMessage ('Deletion stage failed for: %s' % par2head)
 403                         return -EDELETE
 404
 405                 logger.addMessage ('Successfully completed: %s' % par2head, True)
 406                 return SUCCESS
 407
 408 def delete_list (files, interactive=False):
 409         # Delete a list of files
 410
 411         done = False
 412         valid_y = ['Y', 'YES']
 413         valid_n = ['N', 'NO']
 414
 415         if interactive:
 416                 while not done:
 417                         print 'Do you want to delete the following?:'
 418                         s = raw_input ('Delete [y/N]: ').upper()
 419
 420                         if s in valid_y + valid_n:
 421                                 done = True
 422
 423                 if s in valid_n:
 424                         return SUCCESS
 425
 426         for f in files:
 427                 # FIXME: re-enable this in production
 428                 # os.remove (f)
 429                 print 'rm \"%s\"' % f
 430
 431         return SUCCESS
 432
 433
 434 def generate_all_parsets (dir):
 435         # Generate all parsets in the given directory.
 436
 437         assert os.path.isdir (dir) # Directory MUST be valid
 438
 439         parsets = []
 440         p2files = find_all_par2_files (dir)
 441
 442         for f in p2files:
 443                 p = PAR2Set (dir, f)
 444                 if p not in parsets:
 445                         parsets.append (p)
 446
 447         return parsets
 448
 449 def main ():
 450         TOPDIR = os.path.abspath ('test_material')
 451
 452         for (dir, subdirs, files) in os.walk (TOPDIR):
 453                 parsets = generate_all_parsets (dir)
 454                 for p in parsets:
 455                         p.run_all ()
 456
 457         print '\nRARSLAVE STATUS\n'
 458         logger.printAllMessages ()
 459
 460 if __name__ == '__main__':
 461         main ()