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