[RARSLAVE] Add porcelain
[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         options.work_dir = full_abspath (options.work_dir)
483
484         if options.extract_dir != None:
485                 options.extract_dir = full_abspath (options.extract_dir)
486
487         if options.version:
488                 print PROGRAM + ' - ' + VERSION
489                 print
490                 print 'Copyright (c) 2005,2006 Ira W. Snyder (devel@irasnyder.com)'
491                 print
492                 print 'This program comes with ABSOLUTELY NO WARRANTY.'
493                 print 'This is free software, and you are welcome to redistribute it'
494                 print 'under certain conditions. See the file COPYING for details.'
495                 sys.exit (0)
496
497         if options.check_progs:
498                 check_required_progs ()
499
500         if options.write_def_config:
501                 config.write_config (default=True)
502
503         if options.write_config:
504                 config.write_config ()
505
506 def find_loglevel (options):
507
508         loglevel = options.verbose - options.quiet
509
510         if loglevel < RarslaveLogger.MessageType.Fatal:
511                 loglevel = RarslaveLogger.MessageType.Fatal
512
513         if loglevel > RarslaveLogger.MessageType.Debug:
514                 loglevel = RarslaveLogger.MessageType.Debug
515
516         return loglevel
517
518 def printMessageTable (loglevel):
519
520         if logger.hasFatalMessages ():
521                 print '\nFatal Messages\n' + '=' * 80
522                 logger.printLoglevel (RarslaveLogger.MessageType.Fatal)
523
524         if loglevel == RarslaveLogger.MessageType.Fatal:
525                 return
526
527         if logger.hasNormalMessages ():
528                 print '\nNormal Messages\n' + '=' * 80
529                 logger.printLoglevel (RarslaveLogger.MessageType.Normal)
530
531         if loglevel == RarslaveLogger.MessageType.Normal:
532                 return
533
534         if logger.hasVerboseMessages ():
535                 print '\nVerbose Messages\n' + '=' * 80
536                 logger.printLoglevel (RarslaveLogger.MessageType.Verbose)
537
538         if loglevel == RarslaveLogger.MessageType.Verbose:
539                 return
540
541         if logger.hasDebugMessages ():
542                 print '\nDebug Messages\n' + '=' * 80
543                 logger.printLoglevel (RarslaveLogger.MessageType.Debug)
544
545         return
546
547 def main ():
548
549         # Build the OptionParser
550         parser = optparse.OptionParser()
551         parser.add_option('-n', '--not-recursive',
552                                                 action='store_false', dest='recursive',
553                                                 default=config.get_value('options', 'recursive'),
554                                                 help="Don't run recursively")
555
556         parser.add_option('-d', '--work-dir',
557                                                 dest='work_dir', type='string',
558                                                 default=config.get_value('directories', 'working_directory'),
559                                                 help="Start running at DIR", metavar='DIR')
560
561         parser.add_option('-e', '--extract-dir',
562                                                 dest='extract_dir', type='string',
563                                                 default=config.get_value('directories', 'extract_directory'),
564                                                 help="Extract to DIR", metavar='DIR')
565
566         parser.add_option('-p', '--check-required-programs',
567                                                 action='store_true', dest='check_progs',
568                                                 default=False,
569                                                 help="Check for required programs")
570
571         parser.add_option('-f', '--write-default-config',
572                                                 action='store_true', dest='write_def_config',
573                                                 default=False, help="Write out a new default config")
574
575         parser.add_option('-c', '--write-new-config',
576                                                 action='store_true', dest='write_config',
577                                                 default=False, help="Write out the current config")
578
579         parser.add_option('-i', '--interactive', dest='interactive', action='store_true',
580                                                 default=config.get_value('options', 'interactive'),
581                                                 help="Confirm before removing files")
582
583         parser.add_option('-q', '--quiet', dest='quiet', action='count',
584                                                 default=0, help="Output fatal messages only")
585
586         parser.add_option('-v', '--verbose', dest='verbose', action='count',
587                                                 default=0, help="Output extra information")
588
589         parser.add_option('-V', '--version', dest='version', action='store_true',
590                                                 default=False, help="Output version information")
591
592         parser.version = VERSION
593
594         # Parse the given options
595         global options
596         (options, args) = parser.parse_args()
597
598         # Run any special actions that are needed on these options
599         run_options (options)
600
601         # Find the loglevel using the options given 
602         loglevel = find_loglevel (options)
603
604         # Run recursively
605         if options.recursive:
606                 for (dir, subdirs, files) in os.walk (options.work_dir):
607                         parsets = generate_all_parsets (dir)
608                         for p in parsets:
609                                 p.run_all ()
610
611         # Non-recursive
612         else:
613                 parsets = generate_all_parsets (options.work_dir)
614                 for p in parsets:
615                         p.run_all ()
616
617         # Print the results
618         printMessageTable (loglevel)
619
620         # Done!
621         return 0
622
623 if __name__ == '__main__':
624         main ()
625