From 29f202e3d90265226ab9a0ad96f47a4e87d9948c Mon Sep 17 00:00:00 2001
From: maride <git@maride.cc>
Date: Sat, 21 Nov 2020 12:53:07 +0100
Subject: [PATCH] Add FileManager

---
 output/file.go        |  24 ++++++++
 output/filemanager.go | 126 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 150 insertions(+)
 create mode 100644 output/file.go
 create mode 100644 output/filemanager.go

diff --git a/output/file.go b/output/file.go
new file mode 100644
index 0000000..cb1ab3e
--- /dev/null
+++ b/output/file.go
@@ -0,0 +1,24 @@
+package output
+
+import (
+	"crypto/sha256"
+	"fmt"
+)
+
+type File struct {
+	name string
+	content []byte
+	origin string
+	hash string
+}
+
+// Creates a new file object and calculates the hash of the given content
+func NewFile(name string, content []byte, origin string) File {
+	hash := fmt.Sprintf("%x", sha256.Sum256(content))[:6]
+	return File{
+		name:    name,
+		content: content,
+		origin:  origin,
+		hash:    hash,
+	}
+}
\ No newline at end of file
diff --git a/output/filemanager.go b/output/filemanager.go
new file mode 100644
index 0000000..30bc092
--- /dev/null
+++ b/output/filemanager.go
@@ -0,0 +1,126 @@
+package output
+
+import (
+	"fmt"
+	"git.darknebu.la/maride/pancap/common"
+	"io/ioutil"
+	"log"
+	"os"
+	"strings"
+)
+
+var (
+	registeredFiles []File
+	notFound []string
+	extractedFiles int
+)
+
+// Registers a file with the given name and content.
+// This function takes care of filesystem I/O handling and flag parsing.
+// This means that a module should _always_ call this function when a file is encountered.
+// origin is a descriptive string where the file comes from, e.g. the module name.
+func RegisterFile(filename string, content []byte, origin string) {
+	thisFile := NewFile(filename, content, origin)
+	// To avoid doubles, we need to check if that hash is already present
+	for _, f := range registeredFiles {
+		if f.hash == thisFile.hash {
+			// Found - stop here
+			log.Printf("Avoided registering file '%s' because it has the same content as an already registered file ", f.name)
+			return
+		}
+	}
+
+	// None found, add to list
+	registeredFiles = append(registeredFiles, thisFile)
+}
+
+// Iterates over all registered files and checks if they should be extracted and stored, and does exactly that.
+func StoreFiles() {
+	var filesToExtract []File
+
+	// Check different flag scenarios
+	if *targetAllFiles {
+		// We should extract all files.
+		filesToExtract = registeredFiles
+	} else {
+		// We should extract only a given set of files
+		fileList := strings.Split(*targetFiles, ",")
+		for _, f := range fileList {
+			// Iterate over desired files
+			found := false
+			for _, a := range registeredFiles {
+				// Iterate over available (registered) files
+				if f == a.hash {
+					// Found the file
+					found = true
+					filesToExtract = append(filesToExtract, a)
+					break
+				}
+			}
+
+			if !found {
+				// No file found, notify user
+				notFound = append(notFound, fmt.Sprintf("File with hash %s requested but not found.", f))
+			}
+		}
+	}
+
+	// Iterate over all target files and write it them out
+	for _, f := range filesToExtract {
+		writeOut(f)
+	}
+}
+
+// Writes the given file object to disk, along with a stats file placed next to it.
+func writeOut(f File) {
+	targetName := fmt.Sprintf("%s%c%s", *targetOutput, os.PathSeparator, f.hash)
+	targetDescName := fmt.Sprintf("%s.info", targetName)
+	targetDescription := fmt.Sprintf("Filename: %s\nHash: %s\nOrigin: %s\nSize: %d", f.name, f.hash, f.origin, len(f.content))
+
+	// Write target file
+	targetWriteErr := ioutil.WriteFile(targetName, f.content, 0644)
+	if targetWriteErr != nil {
+		log.Printf("Unable to write file %s: %s", targetName, targetWriteErr.Error())
+		return
+	}
+
+	// Write stats file
+	statsWriteErr := ioutil.WriteFile(targetDescName, []byte(targetDescription), 0644)
+	if statsWriteErr != nil {
+		log.Printf("Unable to write file %s: %s", targetName, targetWriteErr.Error())
+		return
+	}
+
+	// Raise stats
+	extractedFiles++
+}
+
+// Prints a brief summary about the extracted files
+func PrintSummary() {
+	summary := fmt.Sprintf("%d files found in stream.\n%d files extracted from stream.", len(registeredFiles), extractedFiles)
+
+	// Generate list of found files
+	var strFileList []string
+	for _, f := range registeredFiles {
+		name := f.name
+		if name == "" {
+			name = "(no name found)"
+		}
+
+		strFileList = append(strFileList, fmt.Sprintf("%s: %s (%s), %d bytes", f.hash, name, f.origin, len(f.content)))
+	}
+
+	// Print list of files as a tree
+	if len(strFileList) > 0 {
+		summary += "\nFound files:"
+		summary += "\n" + common.GenerateTree(strFileList)
+	}
+
+	// Check if we left a few requested files unanswered
+	if len(notFound) > 0 {
+		summary += "\nUnable to find requested file(s) " + strings.Join(notFound, ", ")
+	}
+
+	// Print constructed summary
+	PrintBlock("Files", summary)
+}