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