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