Merge pull request #7776 from gdicristofaro/csvExportIssue

csv export fix
This commit is contained in:
eugene7646 2023-06-09 10:16:15 -04:00 committed by GitHub
commit e7cdf5982d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -18,20 +18,24 @@
*/ */
package org.sleuthkit.autopsy.directorytree; package org.sleuthkit.autopsy.directorytree;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.databind.SequenceWriter;
import com.fasterxml.jackson.dataformat.csv.CsvMapper;
import com.fasterxml.jackson.dataformat.csv.CsvSchema;
import java.awt.Component; import java.awt.Component;
import java.awt.event.ActionEvent; import java.awt.event.ActionEvent;
import java.io.File; import java.io.File;
import java.io.BufferedWriter;
import java.io.FileOutputStream;
import java.io.OutputStreamWriter;
import java.lang.reflect.InvocationTargetException; import java.lang.reflect.InvocationTargetException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Arrays; import java.util.Arrays;
import java.util.Calendar; import java.util.Calendar;
import java.util.Collection; import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator; import java.util.Iterator;
import java.util.List; import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutionException;
import java.util.logging.Level; import java.util.logging.Level;
import javax.swing.AbstractAction; import javax.swing.AbstractAction;
@ -57,7 +61,8 @@ import org.sleuthkit.autopsy.guiutils.JFileChooserFactory;
* Exports CSV version of result nodes to a location selected by the user. * Exports CSV version of result nodes to a location selected by the user.
*/ */
public final class ExportCSVAction extends AbstractAction { public final class ExportCSVAction extends AbstractAction {
// number of rows to sample for different columns
private static final int COLUMN_SAMPLING_ROW_NUM = 100;
private static final Logger logger = Logger.getLogger(ExportCSVAction.class.getName()); private static final Logger logger = Logger.getLogger(ExportCSVAction.class.getName());
private final static String DEFAULT_FILENAME = "Results"; private final static String DEFAULT_FILENAME = "Results";
private final static List<String> columnsToSkip = Arrays.asList(AbstractFilePropertyType.SCORE.toString(), private final static List<String> columnsToSkip = Arrays.asList(AbstractFilePropertyType.SCORE.toString(),
@ -276,43 +281,64 @@ public final class ExportCSVAction extends AbstractAction {
progress.start(); progress.start();
progress.switchToIndeterminate(); progress.switchToIndeterminate();
try (BufferedWriter br = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(outputFile), StandardCharsets.UTF_8))) { if (this.isCancelled()) {
// Write BOM return null;
br.write('\ufeff'); }
// Write the header Set<String> columnHeaderStrs = new HashSet<>();
List<String> headers = new ArrayList<>(); List<CsvSchema.Column> columnHeaders = new ArrayList<>();
PropertySet[] sets = nodesToExport.iterator().next().getPropertySets(); int remainingRowsToSample = 0;
for(PropertySet set : sets) { int columnIdx = 0;
for (Property<?> prop : set.getProperties()) { for (Node nd: nodesToExport) {
if ( ! columnsToSkip.contains(prop.getDisplayName())) { // sample up to 100 rows
headers.add(prop.getDisplayName()); if (remainingRowsToSample >= COLUMN_SAMPLING_ROW_NUM) {
break;
} }
} remainingRowsToSample++;
}
br.write(listToCSV(headers));
for (PropertySet ps: nd.getPropertySets()) {
for (Property prop: ps.getProperties()) {
if (!columnHeaderStrs.contains(prop.getDisplayName()) && !columnsToSkip.contains(prop.getName())) {
columnHeaderStrs.add(prop.getDisplayName());
columnHeaders.add(new CsvSchema.Column(columnIdx, prop.getDisplayName()));
columnIdx++;
}
}
}
}
if (this.isCancelled()) {
return null;
}
CsvSchema schema = CsvSchema.builder()
.addColumns(columnHeaders)
.setUseHeader(true)
.setNullValue("")
.build();
CsvMapper mapper = new CsvMapper();
ObjectWriter writer = mapper.writerFor(Map.class).with(schema);
try (SequenceWriter seqWriter = writer.writeValues(outputFile)) {
// Write each line // Write each line
Iterator<?> nodeIterator = nodesToExport.iterator(); Iterator<?> nodeIterator = nodesToExport.iterator();
while (nodeIterator.hasNext()) { while (nodeIterator.hasNext()) {
if (this.isCancelled()) { if (this.isCancelled()) {
break; return null;
} }
Map<String, Object> rowMap = new HashMap<>();
Node node = (Node)nodeIterator.next(); Node node = (Node)nodeIterator.next();
List<String> values = new ArrayList<>(); for(PropertySet set : node.getPropertySets()) {
sets = node.getPropertySets();
for(PropertySet set : sets) {
for (Property<?> prop : set.getProperties()) { for (Property<?> prop : set.getProperties()) {
if ( ! columnsToSkip.contains(prop.getDisplayName())) { if (!columnsToSkip.contains(prop.getName())) {
values.add(escapeQuotes(prop.getValue().toString())); rowMap.put(prop.getDisplayName(), prop.getValue());
} }
} }
} }
br.write(listToCSV(values)); seqWriter.write(rowMap);
} }
} }
return null; return null;
} }