Copy collector enhanced (#237)

Copy collector enhanced

Author: manavellamnimble

cmd/troubleshoot/cli/run.go

@@ -1,16 +1,19 @@
 package cli
 
 import (
+	"archive/tar"
 	"bytes"
 	"context"
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/url"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"time"
 
@@ -395,7 +398,7 @@ func runCollectors(v *viper.Viper, collectors []*troubleshootv1beta1.Collect, ad
 		}
 
 		if result != nil {
-			err = saveCollectorOutput(result, bundlePath)
+			err = saveCollectorOutput(result, bundlePath, collector)
 			if err != nil {
 				progressChan <- fmt.Errorf("failed to parse collector spec %q: %v", collector.GetDisplayName(), err)
 				continue
@@ -415,8 +418,15 @@ func runCollectors(v *viper.Viper, collectors []*troubleshootv1beta1.Collect, ad
 	return filename, nil
 }
 
-func saveCollectorOutput(output map[string][]byte, bundlePath string) error {
+func saveCollectorOutput(output map[string][]byte, bundlePath string, c *collect.Collector) error {
 	for filename, maybeContents := range output {
+		if c.Collect.Copy != nil {
+			err := untarAndSave(maybeContents, filepath.Join(bundlePath, filepath.Dir(filename)))
+			if err != nil {
+				return errors.Wrap(err, "extract copied files")
+			}
+			continue
+		}
 		fileDir, fileName := filepath.Split(filename)
 		outPath := filepath.Join(bundlePath, fileDir)
 
@@ -431,7 +441,59 @@ func saveCollectorOutput(output map[string][]byte, bundlePath string) error {
 
 	return nil
 }
-
+func untarAndSave(tarFile []byte, bundlePath string) error {
+	keys := make([]string, 0)
+	dirs := make(map[string]*tar.Header)
+	files := make(map[string][]byte)
+	fileHeaders := make(map[string]*tar.Header)
+	tarReader := tar.NewReader(bytes.NewBuffer(tarFile))
+	//Extract and separate tar contentes in file and folders, keeping header info from each one.
+	for {
+		header, err := tarReader.Next()
+		if err != nil {
+			if err != io.EOF {
+				return err
+			}
+			break
+		}
+		switch header.Typeflag {
+		case tar.TypeDir:
+			dirs[header.Name] = header
+		case tar.TypeReg:
+			file := new(bytes.Buffer)
+			_, err = io.Copy(file, tarReader)
+			if err != nil {
+				return err
+			}
+			files[header.Name] = file.Bytes()
+			fileHeaders[header.Name] = header
+		default:
+			return fmt.Errorf("Tar file entry %s contained unsupported file type %v", header.Name, header.FileInfo().Mode())
+		}
+	}
+	//Create directories from base path: <namespace>/<pod name>/containerPath
+	if err := os.MkdirAll(filepath.Join(bundlePath), 0777); err != nil {
+		return errors.Wrap(err, "create output file")
+	}
+	//Order folders stored in variable keys to start always by parent folder. That way folder info is preserved.
+	for k := range dirs {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	//Orderly create folders.
+	for _, k := range keys {
+		if err := os.Mkdir(filepath.Join(bundlePath, k), dirs[k].FileInfo().Mode().Perm()); err != nil {
+			return errors.Wrap(err, "create output file")
+		}
+	}
+	//Populate folders with respective files and its permissions stored in the header.
+	for k, v := range files {
+		if err := ioutil.WriteFile(filepath.Join(bundlePath, k), v, fileHeaders[k].FileInfo().Mode().Perm()); err != nil {
+			return err
+		}
+	}
+	return nil
+}
 func uploadSupportBundle(r *troubleshootv1beta1.ResultRequest, archivePath string) error {
 	contentType := getExpectedContentType(r.URI)
 	if contentType != "" && contentType != "application/tar+gzip" {

pkg/collect/copy.go

@@ -13,6 +13,7 @@ import (
 	"k8s.io/client-go/tools/remotecommand"
 )
 
+//Copy function gets a file or folder from a container specified in the specs.
 func Copy(c *Collector, copyCollector *troubleshootv1beta1.Copy) (map[string][]byte, error) {
 	client, err := kubernetes.NewForConfig(c.ClientConfig)
 	if err != nil {
@@ -47,7 +48,7 @@ func Copy(c *Collector, copyCollector *troubleshootv1beta1.Copy) (map[string][]b
 			}
 
 			for k, v := range files {
-				copyOutput[filepath.Join(bundlePath, k)] = v
+				copyOutput[filepath.Join(bundlePath, filepath.Dir(copyCollector.ContainerPath), k)] = v
 			}
 		}
 	}
@@ -60,9 +61,7 @@ func copyFiles(c *Collector, client *kubernetes.Clientset, pod corev1.Pod, copyC
 	if copyCollector.ContainerName != "" {
 		container = copyCollector.ContainerName
 	}
-
-	command := []string{"cat", copyCollector.ContainerPath}
-
+	command := []string{"tar", "-C", filepath.Dir(copyCollector.ContainerPath), "-cf", "-", filepath.Base(copyCollector.ContainerPath)}
 	req := client.CoreV1().RESTClient().Post().Resource("pods").Name(pod.Name).Namespace(pod.Namespace).SubResource("exec")
 	scheme := runtime.NewScheme()
 	if err := corev1.AddToScheme(scheme); err != nil {
@@ -110,7 +109,7 @@ func copyFiles(c *Collector, client *kubernetes.Clientset, pod corev1.Pod, copyC
 	}
 
 	return map[string][]byte{
-		copyCollector.ContainerPath: output.Bytes(),
+		filepath.Base(copyCollector.ContainerPath) + ".tar": output.Bytes(),
 	}, nil
 }
 

pkg/collect/redact.go

@@ -1,20 +1,129 @@
 package collect
 
 import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"encoding/binary"
+	"io"
+	"path/filepath"
+	"strings"
+
 	troubleshootv1beta1 "github.com/replicatedhq/troubleshoot/pkg/apis/troubleshoot/v1beta1"
 	"github.com/replicatedhq/troubleshoot/pkg/redact"
 )
 
 func redactMap(input map[string][]byte, additionalRedactors []*troubleshootv1beta1.Redact) (map[string][]byte, error) {
 	result := make(map[string][]byte)
 	for k, v := range input {
-		if v != nil {
-			redacted, err := redact.Redact(v, k, additionalRedactors)
+		if v == nil {
+			continue
+		}
+		//If the file is .tar, .tgz or .tar.gz, it must not be redacted. Instead it is decompressed and each file inside the
+		//tar is decompressed, redacted and compressed back into the tar.
+		if filepath.Ext(k) == ".tar" || filepath.Ext(k) == ".tgz" || strings.HasSuffix(k, ".tar.gz") {
+			tarFile := bytes.NewBuffer(v)
+			unRedacted, tarHeaders, err := decompressFile(tarFile, k)
+			if err != nil {
+				return nil, err
+			}
+			redacted, err := redactMap(unRedacted, additionalRedactors)
+			if err != nil {
+				return nil, err
+			}
+			result[k], err = compressFiles(redacted, tarHeaders, k)
 			if err != nil {
 				return nil, err
 			}
-			result[k] = redacted
+			//Content of the tar file was redacted. Continue to next file.
+			continue
 		}
+		redacted, err := redact.Redact(v, k, additionalRedactors)
+		if err != nil {
+			return nil, err
+		}
+		result[k] = redacted
 	}
 	return result, nil
 }
+
+func compressFiles(tarContent map[string][]byte, tarHeaders map[string]*tar.Header, filename string) ([]byte, error) {
+	buff := new(bytes.Buffer)
+	var tw *tar.Writer
+	var zw *gzip.Writer
+	if filepath.Ext(filename) != ".tar" {
+		zw = gzip.NewWriter(buff)
+		tw = tar.NewWriter(zw)
+		defer zw.Close()
+	} else {
+		tw = tar.NewWriter(buff)
+	}
+	defer tw.Close()
+	for p, f := range tarContent {
+		if tarHeaders[p].FileInfo().IsDir() {
+			err := tw.WriteHeader(tarHeaders[p])
+			if err != nil {
+				return nil, err
+			}
+			continue
+		}
+		//File size must be recalculated in case the redactor added some bytes while redacting.
+		tarHeaders[p].Size = int64(binary.Size(f))
+		err := tw.WriteHeader(tarHeaders[p])
+		if err != nil {
+			return nil, err
+		}
+		_, err = tw.Write(f)
+		if err != nil {
+			return nil, err
+		}
+	}
+	err := tw.Close()
+	if err != nil {
+		return nil, err
+	}
+	if filepath.Ext(filename) != ".tar" {
+		err = zw.Close()
+		if err != nil {
+			return nil, err
+		}
+	}
+	return buff.Bytes(), nil
+
+}
+
+func decompressFile(tarFile *bytes.Buffer, filename string) (map[string][]byte, map[string]*tar.Header, error) {
+	var tarReader *tar.Reader
+	var zr *gzip.Reader
+	var err error
+	if filepath.Ext(filename) != ".tar" {
+		zr, err = gzip.NewReader(tarFile)
+		if err != nil {
+			return nil, nil, err
+		}
+		defer zr.Close()
+		tarReader = tar.NewReader(zr)
+	} else {
+		tarReader = tar.NewReader(tarFile)
+	}
+	tarHeaders := make(map[string]*tar.Header)
+	tarContent := make(map[string][]byte)
+	for {
+		header, err := tarReader.Next()
+		if err != nil {
+			if err != io.EOF {
+				return nil, nil, err
+			}
+			break
+		}
+		file := new(bytes.Buffer)
+		_, err = io.Copy(file, tarReader)
+		if err != nil {
+			return nil, nil, err
+		}
+		tarContent[header.Name] = file.Bytes()
+		tarHeaders[header.Name] = header
+
+	}
+	return tarContent, tarHeaders, nil
+}