Skip to content

Commit 4d0cab9

Browse files
committed
in memory fs poc
1 parent 35b4f9d commit 4d0cab9

20 files changed

+611
-176
lines changed

acquisition/acquisition.go

Lines changed: 83 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/mvt-project/androidqf/assets"
2222
"github.com/mvt-project/androidqf/log"
2323
"github.com/mvt-project/androidqf/utils"
24+
"github.com/spf13/afero"
2425
)
2526

2627
// Acquisition is the main object containing all phone information
@@ -35,6 +36,9 @@ type Acquisition struct {
3536
SdCard string `json:"sdcard"`
3637
Cpu string `json:"cpu"`
3738
closeLog func() `json:"-"`
39+
Fs afero.Fs `json:"-"`
40+
UseMemoryFs bool `json:"use_memory_fs"`
41+
KeyFilePresent bool `json:"key_file_present"`
3842
}
3943

4044
// New returns a new Acquisition instance.
@@ -50,21 +54,53 @@ func New(path string) (*Acquisition, error) {
5054
} else {
5155
acq.StoragePath = path
5256
}
53-
// Check if the path exist
54-
stat, err := os.Stat(acq.StoragePath)
55-
if os.IsNotExist(err) {
56-
err := os.Mkdir(acq.StoragePath, 0o755)
57-
if err != nil {
58-
return nil, fmt.Errorf("failed to create acquisition folder: %v", err)
57+
58+
// Check if key.txt is present
59+
keyFilePath := filepath.Join(rt.GetExecutableDirectory(), "key.txt")
60+
if _, err := os.Stat(keyFilePath); err == nil {
61+
acq.KeyFilePresent = true
62+
log.Info("Key file detected, checking available memory for secure processing...")
63+
64+
// Check if we have enough memory (4GB) and key file is present
65+
if utils.HasSufficientMemory() {
66+
acq.UseMemoryFs = true
67+
acq.Fs = afero.NewMemMapFs()
68+
log.Info("Using in-memory filesystem for secure data processing")
69+
70+
// Create the directory structure in memory
71+
err := acq.Fs.MkdirAll(acq.StoragePath, 0o755)
72+
if err != nil {
73+
return nil, fmt.Errorf("failed to create acquisition folder in memory: %v", err)
74+
}
75+
} else {
76+
log.Warning("Insufficient memory for in-memory processing, falling back to disk-based storage")
77+
acq.UseMemoryFs = false
78+
acq.Fs = afero.NewOsFs()
5979
}
6080
} else {
61-
if !stat.IsDir() {
62-
return nil, fmt.Errorf("path exist and is not a folder")
81+
acq.KeyFilePresent = false
82+
acq.UseMemoryFs = false
83+
acq.Fs = afero.NewOsFs()
84+
log.Debug("No key file present, using standard disk-based storage")
85+
}
86+
87+
// For disk-based filesystem, ensure the directory exists
88+
if !acq.UseMemoryFs {
89+
stat, err := os.Stat(acq.StoragePath)
90+
if os.IsNotExist(err) {
91+
err := os.Mkdir(acq.StoragePath, 0o755)
92+
if err != nil {
93+
return nil, fmt.Errorf("failed to create acquisition folder: %v", err)
94+
}
95+
} else {
96+
if !stat.IsDir() {
97+
return nil, fmt.Errorf("path exist and is not a folder")
98+
}
6399
}
64100
}
65101

66102
// Get system information first to get tmp folder
67-
err = acq.GetSystemInformation()
103+
err := acq.GetSystemInformation()
68104
if err != nil {
69105
return nil, err
70106
}
@@ -76,9 +112,20 @@ func New(path string) (*Acquisition, error) {
76112
}
77113
acq.Collector = coll
78114

79-
// Init logging file
115+
// Init logging file - always use OS filesystem for logging to ensure it persists
80116
logPath := filepath.Join(acq.StoragePath, "command.log")
81-
closeLog, err := log.EnableFileLog(log.DEBUG, logPath)
117+
var closeLog func()
118+
if acq.UseMemoryFs {
119+
// For memory filesystem, we still want logs to persist on disk
120+
// Create the directory on disk if it doesn't exist
121+
if _, err := os.Stat(acq.StoragePath); os.IsNotExist(err) {
122+
err := os.MkdirAll(acq.StoragePath, 0o755)
123+
if err != nil {
124+
return nil, fmt.Errorf("failed to create logging directory: %v", err)
125+
}
126+
}
127+
}
128+
closeLog, err = log.EnableFileLog(log.DEBUG, logPath)
82129
if err != nil {
83130
return nil, fmt.Errorf("failed to enable file logging: %v", err)
84131
}
@@ -146,7 +193,7 @@ func (a *Acquisition) GetSystemInformation() error {
146193
func (a *Acquisition) HashFiles() error {
147194
log.Info("Generating list of files hashes...")
148195

149-
csvFile, err := os.Create(filepath.Join(a.StoragePath, "hashes.csv"))
196+
csvFile, err := a.Fs.Create(filepath.Join(a.StoragePath, "hashes.csv"))
150197
if err != nil {
151198
return err
152199
}
@@ -155,20 +202,36 @@ func (a *Acquisition) HashFiles() error {
155202
csvWriter := csv.NewWriter(csvFile)
156203
defer csvWriter.Flush()
157204

158-
_ = filepath.Walk(a.StoragePath, func(filePath string, fileInfo os.FileInfo, err error) error {
205+
err = afero.Walk(a.Fs, a.StoragePath, func(filePath string, fileInfo os.FileInfo, err error) error {
159206
if err != nil {
160207
return err
161208
}
162209

163210
if fileInfo.IsDir() {
164211
return nil
165212
}
166-
// Makes files read only
167-
os.Chmod(filePath, 0o400)
168213

169-
sha256, err := hashes.FileSHA256(filePath)
170-
if err != nil {
171-
return err
214+
// For disk-based filesystem, make files read only
215+
if !a.UseMemoryFs {
216+
os.Chmod(filePath, 0o400)
217+
}
218+
219+
// For memory filesystem, we need to read the file and calculate hash manually
220+
var sha256 string
221+
if a.UseMemoryFs {
222+
content, err := afero.ReadFile(a.Fs, filePath)
223+
if err != nil {
224+
return err
225+
}
226+
sha256, err = hashes.StringSHA256(string(content))
227+
if err != nil {
228+
return err
229+
}
230+
} else {
231+
sha256, err = hashes.FileSHA256(filePath)
232+
if err != nil {
233+
return err
234+
}
172235
}
173236

174237
err = csvWriter.Write([]string{filePath, sha256})
@@ -179,7 +242,7 @@ func (a *Acquisition) HashFiles() error {
179242
return nil
180243
})
181244

182-
return nil
245+
return err
183246
}
184247

185248
func (a *Acquisition) StoreInfo() error {
@@ -193,7 +256,7 @@ func (a *Acquisition) StoreInfo() error {
193256

194257
infoPath := filepath.Join(a.StoragePath, "acquisition.json")
195258

196-
err = os.WriteFile(infoPath, info, 0o644)
259+
err = afero.WriteFile(a.Fs, infoPath, info, 0o644)
197260
if err != nil {
198261
return fmt.Errorf("failed to write acquisition details to file: %v",
199262
err)

acquisition/secure.go

Lines changed: 132 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,66 @@ import (
1616
"filippo.io/age"
1717
saveRuntime "github.com/botherder/go-savetime/runtime"
1818
"github.com/mvt-project/androidqf/log"
19+
"github.com/spf13/afero"
1920
)
2021

21-
func createZipFile(sourceDir, zipPath string) error {
22+
// createZipFromFs creates a zip file from any afero filesystem
23+
func createZipFromFs(fs afero.Fs, sourceDir string, writer io.Writer) error {
24+
zipWriter := zip.NewWriter(writer)
25+
defer zipWriter.Close()
26+
27+
// Walk the filesystem and add all files to the zip
28+
err := afero.Walk(fs, sourceDir, func(path string, info os.FileInfo, err error) error {
29+
if err != nil {
30+
return err
31+
}
32+
33+
// Skip directories
34+
if info.IsDir() {
35+
return nil
36+
}
37+
38+
// Get relative path from sourceDir
39+
relPath, err := filepath.Rel(sourceDir, path)
40+
if err != nil {
41+
return err
42+
}
43+
44+
// Create the file header
45+
header, err := zip.FileInfoHeader(info)
46+
if err != nil {
47+
return err
48+
}
49+
header.Name = filepath.ToSlash(relPath) // Ensure forward slashes for zip compatibility
50+
header.Method = zip.Deflate
51+
52+
// Create the file in the zip
53+
writer, err := zipWriter.CreateHeader(header)
54+
if err != nil {
55+
return err
56+
}
57+
58+
// Read the file content from the filesystem
59+
file, err := fs.Open(path)
60+
if err != nil {
61+
return err
62+
}
63+
defer file.Close()
64+
65+
// Copy file content to zip
66+
_, err = io.Copy(writer, file)
67+
return err
68+
})
69+
70+
if err != nil {
71+
return fmt.Errorf("failed to walk filesystem: %v", err)
72+
}
73+
74+
return zipWriter.Close()
75+
}
76+
77+
// createZipFromDisk creates a zip file from disk (legacy approach)
78+
func createZipFromDisk(sourceDir, zipPath string) error {
2279
zipFile, err := os.Create(zipPath)
2380
if err != nil {
2481
return fmt.Errorf("failed to create ZIP file: %v", err)
@@ -39,83 +96,111 @@ func createZipFile(sourceDir, zipPath string) error {
3996
}
4097

4198
func (a *Acquisition) StoreSecurely() error {
42-
cwd := saveRuntime.GetExecutableDirectory()
43-
44-
keyFilePath := filepath.Join(cwd, "key.txt")
45-
if _, err := os.Stat(keyFilePath); os.IsNotExist(err) {
99+
// Only proceed if key file is present
100+
if !a.KeyFilePresent {
46101
return nil
47102
}
48103

49-
log.Info("You provided an age public key, storing the acquisition securely.")
50-
51-
zipFileName := fmt.Sprintf("%s.zip", a.UUID)
52-
zipFilePath := filepath.Join(cwd, zipFileName)
104+
log.Info("Age public key detected, storing the acquisition securely.")
53105

54-
log.Info("Compressing the acquisition folder. This might take a while...")
55-
56-
err := createZipFile(a.StoragePath, zipFilePath)
57-
if err != nil {
58-
return err
59-
}
60-
61-
log.Info("Encrypting the compressed archive. This might take a while...")
106+
cwd := saveRuntime.GetExecutableDirectory()
107+
keyFilePath := filepath.Join(cwd, "key.txt")
62108

109+
// Read the public key
63110
publicKey, err := os.ReadFile(keyFilePath)
64111
if err != nil {
65-
return err
112+
return fmt.Errorf("failed to read key file: %v", err)
66113
}
67114
publicKeyStr := strings.TrimSpace(string(publicKey))
68115

116+
// Parse the age recipient
69117
recipient, err := age.ParseX25519Recipient(publicKeyStr)
70118
if err != nil {
71119
return fmt.Errorf("failed to parse public key %q: %v", publicKeyStr, err)
72120
}
73121

74-
zipFile, err := os.Open(zipFilePath)
75-
if err != nil {
76-
return err
77-
}
78-
defer zipFile.Close()
79-
80-
encFileName := fmt.Sprintf("%s.age", zipFileName)
122+
// Create the encrypted file
123+
encFileName := fmt.Sprintf("%s.age", a.UUID)
81124
encFilePath := filepath.Join(cwd, encFileName)
82-
encFile, err := os.OpenFile(encFilePath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0o600)
125+
encFile, err := os.OpenFile(encFilePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
83126
if err != nil {
84127
return fmt.Errorf("unable to create encrypted file: %v", err)
85128
}
86129
defer encFile.Close()
87130

88-
w, err := age.Encrypt(encFile, recipient)
131+
// Create age encryptor
132+
ageWriter, err := age.Encrypt(encFile, recipient)
89133
if err != nil {
90-
return fmt.Errorf("failed to create encrypted file: %v", err)
134+
return fmt.Errorf("failed to create age encryptor: %v", err)
91135
}
92136

93-
_, err = io.Copy(w, zipFile)
94-
if err != nil {
95-
return fmt.Errorf("failed to write to encrypted file: %v", err)
137+
if a.UseMemoryFs {
138+
log.Info("Compressing and encrypting acquisition data from memory. This might take a while...")
139+
140+
// Create zip directly from memory filesystem and encrypt it
141+
err = createZipFromFs(a.Fs, a.StoragePath, ageWriter)
142+
if err != nil {
143+
return fmt.Errorf("failed to create encrypted zip from memory: %v", err)
144+
}
145+
} else {
146+
log.Info("Compressing the acquisition folder. This might take a while...")
147+
148+
// Create temporary zip file on disk first
149+
zipFileName := fmt.Sprintf("%s.zip", a.UUID)
150+
zipFilePath := filepath.Join(cwd, zipFileName)
151+
152+
err := createZipFromDisk(a.StoragePath, zipFilePath)
153+
if err != nil {
154+
return fmt.Errorf("failed to create zip file: %v", err)
155+
}
156+
157+
log.Info("Encrypting the compressed archive. This might take a while...")
158+
159+
// Read the zip file and encrypt it
160+
zipFile, err := os.Open(zipFilePath)
161+
if err != nil {
162+
return fmt.Errorf("failed to open zip file: %v", err)
163+
}
164+
165+
_, err = io.Copy(ageWriter, zipFile)
166+
zipFile.Close()
167+
168+
if err != nil {
169+
return fmt.Errorf("failed to encrypt zip file: %v", err)
170+
}
171+
172+
// Remove the temporary unencrypted zip file
173+
err = os.Remove(zipFilePath)
174+
if err != nil {
175+
log.Warningf("Failed to remove temporary zip file: %v", err)
176+
}
96177
}
97178

98-
if err := w.Close(); err != nil {
179+
// Close the age writer
180+
if err := ageWriter.Close(); err != nil {
99181
return fmt.Errorf("failed to close encrypted file: %v", err)
100182
}
101183

102184
log.Infof("Acquisition successfully encrypted at %s", encFilePath)
103185

104-
// TODO: we should securely wipe the files.
105-
zipFile.Close()
106-
err = os.Remove(zipFilePath)
107-
if err != nil {
108-
return fmt.Errorf("failed to delete the unencrypted compressed archive: %v", err)
109-
}
110-
111-
// Ensure log file is closed before removing the acquisition directory
112-
if a.closeLog != nil {
113-
defer a.closeLog()
114-
}
115-
116-
err = os.RemoveAll(a.StoragePath)
117-
if err != nil {
118-
return fmt.Errorf("failed to delete the original unencrypted acquisition folder: %v", err)
186+
// Clean up based on filesystem type
187+
if a.UseMemoryFs {
188+
log.Info("Clearing in-memory data...")
189+
// For memory filesystem, we just need to ensure the log is closed
190+
// The memory will be freed when the filesystem is dereferenced
191+
} else {
192+
log.Info("Removing unencrypted acquisition folder...")
193+
// Ensure log file is closed before removing the acquisition directory
194+
if a.closeLog != nil {
195+
a.closeLog()
196+
a.closeLog = nil // Prevent double-close
197+
}
198+
199+
// Remove the original unencrypted folder
200+
err = os.RemoveAll(a.StoragePath)
201+
if err != nil {
202+
return fmt.Errorf("failed to delete the original unencrypted acquisition folder: %v", err)
203+
}
119204
}
120205

121206
return nil

0 commit comments

Comments
 (0)