-
Notifications
You must be signed in to change notification settings - Fork 0
/
filemanager.go
133 lines (114 loc) · 3.76 KB
/
filemanager.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
package output
import (
"fmt"
"io/ioutil"
"log"
"os"
"strings"
"github.com/maride/pancap/common"
)
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) {
// Check if there even is anything to register
if len(content) == 0 {
// File is empty, won't register the void
log.Printf("Avoided registering file from %s because it is empty.", origin)
return
}
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 from %s because it has the same content as an already registered file ", origin)
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)
}