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