ceb7eac3089a48aea03e9600e2ade12261c0169f
[rarslave2.git] / rarslave.py
1 #!/usr/bin/env python
2 # vim: set ts=4 sts=4 sw=4 textwidth=112 :
3
4 VERSION="2.0.0"
5 PROGRAM="rarslave2"
6
7 import re, os, sys, optparse
8 import par2parser
9 import RarslaveConfig
10 import RarslaveLogger
11
12 # Global Variables
13 (TYPE_OLDRAR, TYPE_NEWRAR, TYPE_ZIP, TYPE_NOEXTRACT) = range (4)
14 (SUCCESS, ECHECK, EEXTRACT, EDELETE) = range(4)
15 config = RarslaveConfig.RarslaveConfig()
16 logger = RarslaveLogger.RarslaveLogger ()
17
18 # Global options to be set / used later.
19 options = None
20
21 class RarslaveExtractor (object):
22
23         def __init__ (self, type):
24                 self.type = type
25                 self.heads = []
26
27         def addHead (self, dir, head):
28                 assert os.path.isdir (dir)
29                 # REQUIRES that the dir is valid, but not that the file is valid, so that
30                 # we can move a file that doesn't exist yet.
31                 # FIXME: probably CAN add this back, since we should be running this AFTER repair.
32                 #assert os.path.isfile (os.path.join (dir, head))
33
34                 full_head = os.path.join (dir, head)
35                 logger.addMessage ('Adding extraction head: %s' % full_head, RarslaveLogger.MessageType.Debug)
36                 self.heads.append (full_head)
37
38         def extract (self, todir=None):
39                 # Extract all heads of this set
40
41                 # Create the directory $todir if it doesn't exist
42                 if todir != None and not os.path.isdir (todir):
43                         logger.addMessage ('Creating directory: %s' % todir, RarslaveLogger.MessageType.Verbose)
44                         try:
45                                 os.makedirs (todir)
46                         except OSError:
47                                 logger.addMessage ('FAILED to create directory: %s' % todir, RarslaveLogger.MessageType.Fatal)
48                                 return -EEXTRACT
49
50                 # Extract all heads
51                 extraction_func = \
52                         { TYPE_OLDRAR : self.__extract_rar,
53                           TYPE_NEWRAR : self.__extract_rar,
54                           TYPE_ZIP    : self.__extract_zip,
55                           TYPE_NOEXTRACT : self.__extract_noextract }[self.type]
56
57                 # Call the extraction function on each head
58                 for h in self.heads:
59                         if todir == None:
60                                 # Run in the head's directory
61                                 ret = extraction_func (h, os.path.dirname (h))
62                         else:
63                                 ret = extraction_func (h, todir)
64
65                         logger.addMessage ('Extraction Function returned: %d' % ret, RarslaveLogger.MessageType.Debug)
66
67                         # Check error code
68                         if ret != SUCCESS:
69                                 logger.addMessage ('Failed extracting: %s' % h, RarslaveLogger.MessageType.Fatal)
70                                 return -EEXTRACT
71
72                 return SUCCESS
73
74         def __extract_rar (self, file, todir):
75                 assert os.path.isfile (file)
76                 assert os.path.isdir (todir)
77
78                 RAR_CMD = config.get_value ('commands', 'unrar')
79
80                 cmd = '%s \"%s\"' % (RAR_CMD, file)
81                 ret = run_command (cmd, todir)
82
83                 # Check error code
84                 if ret != 0:
85                         return -EEXTRACT
86
87                 return SUCCESS
88
89         def __extract_zip (self, file, todir):
90                 ZIP_CMD = config.get_value ('commands', 'unzip')
91
92                 cmd = ZIP_CMD % (file, todir)
93                 ret = run_command (cmd)
94
95                 # Check error code
96                 if ret != 0:
97                         return -EEXTRACT
98
99                 return SUCCESS
100
101         def __extract_noextract (self, file, todir):
102                 # Just move this file to the $todir, since no extraction is needed
103                 # FIXME: NOTE: mv will fail by itself if you're moving to the same dir!
104                 NOEXTRACT_CMD = config.get_value ('commands', 'noextract')
105
106                 cmd = NOEXTRACT_CMD % (file, todir)
107                 ret = run_command (cmd)
108
109                 # Check error code
110                 if ret != 0:
111                         return -EEXTRACT
112
113                 return SUCCESS
114
115
116
117 class RarslaveRepairer (object):
118         # Verify (and repair) the set
119         # Make sure it worked, otherwise clean up and return failure
120
121         def __init__ (self, dir, file, join=False):
122                 self.dir  = dir  # the directory containing the par2 file
123                 self.file = file # the par2 file
124                 self.join = join # True if the par2 set is 001 002 ...
125
126                 assert os.path.isdir (dir)
127                 assert os.path.isfile (os.path.join (dir, file))
128
129         def checkAndRepair (self):
130                 # Form the command:
131                 # par2repair -- PAR2 PAR2_EXTRA [JOIN_FILES]
132                 PAR2_CMD = config.get_value ('commands', 'par2repair')
133
134                 # Get set up
135                 basename = get_basename (self.file)
136                 all_files = find_likely_files (basename, self.dir)
137                 all_files.sort ()
138                 par2_files = find_par2_files (all_files)
139
140                 # assemble the command
141                 command = "%s \"%s\" " % (PAR2_CMD, self.file)
142
143                 for f in par2_files:
144                         if f != self.file:
145                                 command += "\"%s\" " % os.path.split (f)[1]
146
147                 if self.join:
148                         for f in all_files:
149                                 if f not in par2_files:
150                                         command += "\"%s\" " % os.path.split (f)[1]
151
152                 # run the command
153                 ret = run_command (command, self.dir)
154
155                 # check the result
156                 if ret != 0:
157                         logger.addMessage ('PAR2 Check / Repair failed: %s' % self.file, RarslaveLogger.MessageType.Fatal)
158                         return -ECHECK
159
160                 return SUCCESS
161
162 def run_command (cmd, indir=None):
163         # Runs the specified command-line in the directory given (or, in the current directory
164         # if none is given). It returns the status code given by the application.
165
166         pwd = os.getcwd ()
167
168         if indir != None:
169                 assert os.path.isdir (indir) # MUST be a directory!
170                 os.chdir (indir)
171
172         # FIXME: re-enable this after testing
173         print 'RUNNING (%s): %s' % (indir, cmd)
174         return SUCCESS
175
176         # return os.system (cmd)
177
178
179 def full_abspath (p):
180         return os.path.abspath (os.path.expanduser (p))
181
182 def get_basename (name):
183         """Strips most kinds of endings from a filename"""
184
185         regex = config.get_value ('regular expressions', 'basename_regex')
186         r = re.compile (regex, re.IGNORECASE)
187         done = False
188
189         while not done:
190                 done = True
191
192                 if r.match (name):
193                         g = r.match (name).groups()
194                         name = g[0]
195                         done = False
196
197         return name
198
199 def find_likely_files (name, dir):
200         """Finds files which are likely to be part of the set corresponding
201            to $name in the directory $dir"""
202
203         if not os.path.isdir (os.path.abspath (dir)):
204                 raise ValueError # bad directory given
205
206         dir = os.path.abspath (dir)
207         ename = re.escape (name)
208         regex = re.compile ('^%s.*$' % (ename, ))
209
210         return [f for f in os.listdir (dir) if regex.match (f)]
211
212 def find_par2_files (files):
213         """Find all par2 files in the list $files"""
214
215         PAR2_REGEX = config.get_value ('regular expressions', 'par2_regex')
216         regex = re.compile (PAR2_REGEX, re.IGNORECASE)
217         return [f for f in files if regex.match (f)]
218
219 def find_all_par2_files (dir):
220         """Finds all par2 files in a directory"""
221         # NOTE: does NOT return absolute paths
222
223         if not os.path.isdir (os.path.abspath (dir)):
224                 raise ValueError # bad directory given
225
226         dir = os.path.abspath (dir)
227         files = os.listdir (dir)
228
229         return find_par2_files (files)
230
231 def has_extension (f, ext):
232         """Checks if f has the extension ext"""
233
234         if ext[0] != '.':
235                 ext = '.' + ext
236
237         ext = re.escape (ext)
238         regex = re.compile ('^.*%s$' % (ext, ), re.IGNORECASE)
239         return regex.match (f)
240
241 def find_extraction_heads (dir, files):
242         """Takes a list of possible files and finds likely heads of
243            extraction."""
244
245         # NOTE: perhaps this should happen AFTER repair is
246         # NOTE: successful. That way all files would already exist
247
248         # According to various sources online:
249         # 1) pre rar-3.0: .rar .r00 .r01 ...
250         # 2) post rar-3.0: .part01.rar .part02.rar 
251         # 3) zip all ver: .zip 
252
253         extractor = None
254         p2files = find_par2_files (files)
255
256         # Old RAR type, find all files ending in .rar
257         if is_oldrar (files):
258                 extractor = RarslaveExtractor (TYPE_OLDRAR)
259                 regex = re.compile ('^.*\.rar$', re.IGNORECASE)
260                 for f in files:
261                         if regex.match (f):
262                                 extractor.addHead (dir, f)
263
264         if is_newrar (files):
265                 extractor = RarslaveExtractor (TYPE_NEWRAR)
266                 regex = re.compile ('^.*\.part01.rar$', re.IGNORECASE)
267                 for f in files:
268                         if regex.match (f):
269                                 extractor.addHead (dir, f)
270
271         if is_zip (files):
272                 extractor = RarslaveExtractor (TYPE_ZIP)
273                 regex = re.compile ('^.*\.zip$', re.IGNORECASE)
274                 for f in files:
275                         if regex.match (f):
276                                 extractor.addHead (dir, f)
277
278         if is_noextract (files):
279                 # Use the Par2 Parser (from cfv) here to find out what files are protected.
280                 # Since these are not being extracted, they will be mv'd to another directory
281                 # later.
282                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
283
284                 for f in p2files:
285                         done = False
286                         try:
287                                 prot_files = par2parser.get_protected_files (dir, f)
288                                 done = True
289                         except: #FIXME: add the actual exceptions
290                                 logger.addMessage ('Error parsing PAR2 file: %s', f)
291                                 continue
292
293                         if done:
294                                 break
295
296                 if done:
297                         for f in prot_files:
298                                 extractor.addHead (dir, f)
299                 else:
300                         logger.addMessage ('Error parsing all PAR2 files in this set ...')
301
302         # Make sure we found the type
303         if extractor == None:
304                 logger.addMessage ('Not able to find an extractor for this type of set: %s' % p2files[0],
305                                 RarslaveLogger.MessageType.Fatal)
306
307                 # No-heads here, but it's better than failing completely
308                 extractor = RarslaveExtractor (TYPE_NOEXTRACT)
309
310         return extractor
311
312 def is_oldrar (files):
313         for f in files:
314                 if has_extension (f, '.r00'):
315                         return True
316
317         return False
318
319 def is_newrar (files):
320         for f in files:
321                 if has_extension (f, '.part01.rar'):
322                         return True
323
324         return False
325
326 def is_zip (files):
327         for f in files:
328                 if has_extension (f, '.zip'):
329                         return True
330
331         return False
332
333 def is_noextract (files):
334         # Type that needs no extraction.
335         # TODO: Add others ???
336         for f in files:
337                 if has_extension (f, '.001'):
338                         return True
339
340         return False
341
342 def find_deleteable_files (files):
343         # Deleteable types regex should come from the config
344         dfiles = []
345         DELETE_REGEX = config.get_value ('regular expressions', 'delete_regex')
346         dregex = re.compile (DELETE_REGEX, re.IGNORECASE)
347
348         return [f for f in files if dregex.match (f)]
349
350 def printlist (li):
351         for f in li:
352                 print f
353
354 class PAR2Set (object):
355
356         dir = None
357         file = None
358         likely_files = []
359
360         def __init__ (self, dir, file):
361                 assert os.path.isdir (dir)
362                 assert os.path.isfile (os.path.join (dir, file))
363
364                 self.dir = dir
365                 self.file = file
366
367                 basename = get_basename (file)
368                 self.likely_files = find_likely_files (basename, dir)
369
370         def __list_eq (self, l1, l2):
371
372                 if len(l1) != len(l2):
373                         return False
374
375                 for e in l1:
376                         if e not in l2:
377                                 return False
378
379                 return True
380
381         def __eq__ (self, rhs):
382                 return self.__list_eq (self.likely_files, rhs.likely_files)
383
384         def run_all (self):
385                 par2files = find_par2_files (self.likely_files)
386                 par2head = par2files[0]
387
388                 join = is_noextract (self.likely_files)
389
390                 # Repair Stage
391                 repairer = RarslaveRepairer (self.dir, par2head, join)
392                 ret = repairer.checkAndRepair ()
393
394                 if ret != SUCCESS:
395                         logger.addMessage ('Repair stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
396                         return -ECHECK
397
398                 # Extraction Stage
399                 EXTRACT_DIR = options.extract_dir
400                 extractor = find_extraction_heads (self.dir, self.likely_files)
401                 ret = extractor.extract (EXTRACT_DIR)
402
403                 if ret != SUCCESS:
404                         logger.addMessage ('Extraction stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
405                         return -EEXTRACT
406
407                 # Deletion Stage
408                 DELETE_INTERACTIVE = options.interactive
409                 deleteable_files = find_deleteable_files (self.likely_files)
410                 ret = delete_list (deleteable_files, DELETE_INTERACTIVE)
411
412                 if ret != SUCCESS:
413                         logger.addMessage ('Deletion stage failed for: %s' % par2head, RarslaveLogger.MessageType.Fatal)
414                         return -EDELETE
415
416                 logger.addMessage ('Successfully completed: %s' % par2head)
417                 return SUCCESS
418
419 def delete_list (files, interactive=False):
420         # Delete a list of files
421
422         done = False
423         valid_y = ['Y', 'YES']
424         valid_n = ['N', 'NO']
425
426         if interactive:
427                 while not done:
428                         print 'Do you want to delete the following?:'
429                         s = raw_input ('Delete [y/N]: ').upper()
430
431                         if s in valid_y + valid_n:
432                                 done = True
433
434                 if s in valid_n:
435                         return SUCCESS
436
437         for f in files:
438                 # FIXME: re-enable this in production
439                 # os.remove (f)
440                 print 'rm \"%s\"' % f
441
442         return SUCCESS
443
444
445 def generate_all_parsets (dir):
446         # Generate all parsets in the given directory.
447
448         assert os.path.isdir (dir) # Directory MUST be valid
449
450         parsets = []
451         p2files = find_all_par2_files (dir)
452
453         for f in p2files:
454                 p = PAR2Set (dir, f)
455                 if p not in parsets:
456                         parsets.append (p)
457
458         return parsets
459
460 def check_required_progs():
461         """Check if the required programs are installed"""
462
463         shell_not_found = 32512
464         needed = []
465
466         if run_command ('par2repair --help > /dev/null 2>&1') == shell_not_found:
467                 needed.append ('par2repair')
468
469         if run_command ('unrar --help > /dev/null 2>&1') == shell_not_found:
470                 needed.append ('unrar')
471
472         if run_command ('unzip --help > /dev/null 2>&1') == shell_not_found:
473                 needed.append ('unzip')
474
475         if needed:
476                 for n in needed:
477                         print 'Needed program "%s" not found in $PATH' % (n, )
478
479                 sys.exit(1)
480
481 def run_options (options):
482
483         # Fix directories
484         options.work_dir = full_abspath (options.work_dir)
485
486         # Make sure that the directory is valid
487         if not os.path.isdir (options.work_dir):
488                 sys.stderr.write ('\"%s\" is not a valid directory. Use the \"-d\"\n' % options.work_dir)
489                 sys.stderr.write ('option to override the working directory temporarily, or edit the\n')
490                 sys.stderr.write ('configuration file to override the working directory permanently.\n')
491                 sys.exit (1)
492
493         if options.extract_dir != None:
494                 options.extract_dir = full_abspath (options.extract_dir)
495
496         if options.version:
497                 print PROGRAM + ' - ' + VERSION
498                 print
499                 print 'Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)'
500                 print
501                 print 'This program comes with ABSOLUTELY NO WARRANTY.'
502                 print 'This is free software, and you are welcome to redistribute it'
503                 print 'under certain conditions. See the file COPYING for details.'
504                 sys.exit (0)
505
506         if options.check_progs:
507                 check_required_progs ()
508
509         if options.write_def_config:
510                 config.write_config (default=True)
511
512         if options.write_config:
513                 config.write_config ()
514
515 def find_loglevel (options):
516
517         loglevel = options.verbose - options.quiet
518
519         if loglevel < RarslaveLogger.MessageType.Fatal:
520                 loglevel = RarslaveLogger.MessageType.Fatal
521
522         if loglevel > RarslaveLogger.MessageType.Debug:
523                 loglevel = RarslaveLogger.MessageType.Debug
524
525         return loglevel
526
527 def printMessageTable (loglevel):
528
529         if logger.hasFatalMessages ():
530                 print '\nFatal Messages\n' + '=' * 80
531                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
532
533         if loglevel == RarslaveLogger.MessageType.Fatal:
534                 return
535
536         if logger.hasNormalMessages ():
537                 print '\nNormal Messages\n' + '=' * 80
538                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
539
540         if loglevel == RarslaveLogger.MessageType.Normal:
541                 return
542
543         if logger.hasVerboseMessages ():
544                 print '\nVerbose Messages\n' + '=' * 80
545                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
546
547         if loglevel == RarslaveLogger.MessageType.Verbose:
548                 return
549
550         if logger.hasDebugMessages ():
551                 print '\nDebug Messages\n' + '=' * 80
552                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
553
554         return
555
556 def main ():
557
558         # Build the OptionParser
559         parser = optparse.OptionParser()
560         parser.add_option('-n', '--not-recursive',
561                                                 action='store_false', dest='recursive',
562                                                 default=config.get_value('options', 'recursive'),
563                                                 help="Don't run recursively")
564
565         parser.add_option('-d', '--work-dir',
566                                                 dest='work_dir', type='string',
567                                                 default=config.get_value('directories', 'working_directory'),
568                                                 help="Start running at DIR", metavar='DIR')
569
570         parser.add_option('-e', '--extract-dir',
571                                                 dest='extract_dir', type='string',
572                                                 default=config.get_value('directories', 'extract_directory'),
573                                                 help="Extract to DIR", metavar='DIR')
574
575         parser.add_option('-p', '--check-required-programs',
576                                                 action='store_true', dest='check_progs',
577                                                 default=False,
578                                                 help="Check for required programs")
579
580         parser.add_option('-f', '--write-default-config',
581                                                 action='store_true', dest='write_def_config',
582                                                 default=False, help="Write out a new default config")
583
584         parser.add_option('-c', '--write-new-config',
585                                                 action='store_true', dest='write_config',
586                                                 default=False, help="Write out the current config")
587
588         parser.add_option('-i', '--interactive', dest='interactive', action='store_true',
589                                                 default=config.get_value('options', 'interactive'),
590                                                 help="Confirm before removing files")
591
592         parser.add_option('-q', '--quiet', dest='quiet', action='count',
593                                                 default=0, help="Output fatal messages only")
594
595         parser.add_option('-v', '--verbose', dest='verbose', action='count',
596                                                 default=0, help="Output extra information")
597
598         parser.add_option('-V', '--version', dest='version', action='store_true',
599                                                 default=False, help="Output version information")
600
601         parser.version = VERSION
602
603         # Parse the given options
604         global options
605         (options, args) = parser.parse_args()
606
607         # Run any special actions that are needed on these options
608         run_options (options)
609
610         # Find the loglevel using the options given 
611         loglevel = find_loglevel (options)
612
613         # Run recursively
614         if options.recursive:
615                 for (dir, subdirs, files) in os.walk (options.work_dir):
616                         parsets = generate_all_parsets (dir)
617                         for p in parsets:
618                                 p.run_all ()
619
620         # Non-recursive
621         else:
622                 parsets = generate_all_parsets (options.work_dir)
623                 for p in parsets:
624                         p.run_all ()
625
626         # Print the results
627         printMessageTable (loglevel)
628
629         # Done!
630         return 0
631
632 if __name__ == '__main__':
633         main ()
634