Add support for tar.bz2 and tar.gz archives

2 files changed, 18 insertions, 3 deletions

diff --git a/lib/archive.py b/lib/archive.py
index f7649d5..89425cf 100644
--- a/lib/archive.py
+++ b/lib/archive.py
@@ -5,13 +5,18 @@ import mat
 import shutil
 
 class TarStripper(parser.Generic_parser):
+    def compression_type(self):
+        self.compression = ''
+
     def remove_all(self):
+        self.compression_type()
         if not tarfile.is_tarfile(self.filename):
             print('%s is not a valid tar file' % self.filename)
             sys.exit(1)
 
-        tarin = tarfile.open(self.filename, 'r')
-        tarout = tarfile.open(self.filename + parser.POSTFIX, 'w')
+        tarin = tarfile.open(self.filename, 'r' + self.compression)
+        tarout = tarfile.open(self.filename + parser.POSTFIX,
+            'w' + self.compression)
         folder_list = []
 
         for current_file in tarin.getmembers():
@@ -36,3 +41,11 @@ class TarStripper(parser.Generic_parser):
 
     def is_clean(self):
         return False
+
+class GzipStripper(TarStripper):
+    def compression_type(self):
+        self.compression = ':gz'
+
+class Bzip2Stripper(TarStripper):
+    def compression_type(self):
+        self.compression = ':bz2'
diff --git a/lib/mat.py b/lib/mat.py
index 732dc25..9624511 100644
--- a/lib/mat.py
+++ b/lib/mat.py
@@ -25,6 +25,8 @@ strippers = {
     hachoir_parser.audio.MpegAudioFile: audio.MpegAudioStripper,
     hachoir_parser.misc.PDFDocument: misc.PdfStripper,
     hachoir_parser.archive.TarFile: archive.TarStripper,
+    hachoir_parser.archive.gzip_parser.GzipParser: archive.GzipStripper,
+    hachoir_parser.archive.bzip2_parser.Bzip2Parser: archive.Bzip2Stripper,
 }
 
 def is_secure(filename):
@@ -33,7 +35,7 @@ def is_secure(filename):
     '''
 
     if not(os.path.isfile(filename)): #check if the file exist
-        print("Error: %s is not a valid file" % name)
+        print("Error: %s is not a valid file" % filename)
         sys.exit(1)
 
 def create_class_file(name, backup):