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