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