1 files changed, 22 insertions, 8 deletions
diff --git a/archive/unarchive.go b/archive/unarchive.go
index 5089370..d283d4e 100644
--- a/archive/unarchive.go
+++ b/archive/unarchive.go
@@ -4,29 +4,43 @@ import (
 	"archive/tar"
 	fp "path/filepath"
 
+	"compress/gzip"
+
 	"io"
 	"os"
 )
 
 type Unarchive struct {
 	Input     string
-	inputFile *os.File
 	tarReader *tar.Reader
 }
 
 // Extracts a tar arcive to the current working directory
 // This will overwrite everything. So be careful
-func (u *Unarchive) UnTar() error {
-	var err error
-	u.inputFile, err = os.Open(u.Input)
+func (u *Unarchive) UnTar(compressionType string) error {
+	f, err := os.Open(u.Input)
 	if err != nil {
 		return err
 	}
-	defer u.inputFile.Close()
-	u.tarReader = tar.NewReader(u.inputFile)
+	defer f.Close()
+
+	var treader *tar.Reader
+
+	switch compressionType {
+	case "gz":
+		gzreader, err := gzip.NewReader(f)
+		if err != nil {
+			return err
+		}
+		// Read from the gzip reader instead of the file
+		treader = tar.NewReader(gzreader)
+	default:
+		// Read from the file directly
+		treader = tar.NewReader(f)
+	}
 
 	for {
-		hdr, err := u.tarReader.Next()
+		hdr, err := treader.Next()
 		if err == io.EOF {
 			// We've reached the end! Whoee
 			break
@@ -54,7 +68,7 @@ func (u *Unarchive) UnTar() error {
 			return err
 		}
 		defer f.Close()
-		if _, err := io.Copy(f, u.tarReader); err != nil {
+		if _, err := io.Copy(f, treader); err != nil {
 			return err
 		}