ExportCSVAction.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2019-2021 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this content except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
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.event.ActionEvent;
import java.io.File;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.logging.Level;
import javax.swing.AbstractAction;
import javax.swing.JFileChooser;
import javax.swing.JOptionPane;
import javax.swing.SwingWorker;
import javax.swing.filechooser.FileNameExtensionFilter;
import org.netbeans.api.progress.ProgressHandle;
import org.openide.util.Cancellable;
import org.openide.util.NbBundle;
import org.openide.util.Utilities;
import org.sleuthkit.autopsy.casemodule.Case;
import org.sleuthkit.autopsy.casemodule.NoCurrentCaseException;
import org.sleuthkit.autopsy.coreutils.Logger;
import org.sleuthkit.autopsy.coreutils.MessageNotifyUtil;
import org.sleuthkit.autopsy.datamodel.AbstractAbstractFileNode.AbstractFilePropertyType;
import org.openide.nodes.Node;
import org.openide.nodes.Node.PropertySet;
import org.openide.nodes.Node.Property;
import org.sleuthkit.autopsy.guiutils.JFileChooserFactory;

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 final static String DEFAULT_FILENAME = "Results";
    private final static List<String> columnsToSkip = Arrays.asList(AbstractFilePropertyType.SCORE.toString(),
            AbstractFilePropertyType.COMMENT.toString(), AbstractFilePropertyType.OCCURRENCES.toString());

    private static String userDefinedExportPath;

    // This class is a singleton to support multi-selection of nodes, since 
    // org.openide.nodes.NodeOp.findActions(Node[] nodes) will only pick up an Action if every 
    // node in the array returns a reference to the same action object from Node.getActions(boolean).    
    private static ExportCSVAction instance;
    
    private static final JFileChooserFactory chooserHelper = new JFileChooserFactory();

    public static synchronized ExportCSVAction getInstance() {
        if (null == instance) {
            instance = new ExportCSVAction();
        }
        return instance;
    }

    @NbBundle.Messages({"ExportCSV.title.text=Export Selected Rows to CSV"})
    private ExportCSVAction() {
        super(Bundle.ExportCSV_title_text());
    }
    
    @Override
    public void actionPerformed(ActionEvent e) {
        Collection<? extends Node> selectedNodes = Utilities.actionsGlobalContext().lookupAll(Node.class);
        saveNodesToCSV(selectedNodes, (Component)e.getSource());
    }

    @NbBundle.Messages({
        "# {0} - Output file",
        "ExportCSV.saveNodesToCSV.fileExists=File {0} already exists",
        "ExportCSV.saveNodesToCSV.noCurrentCase=No open case available",
        "ExportCSV.saveNodesToCSV.empty=No data to export"})
    public static void saveNodesToCSV(Collection<? extends Node> nodesToExport, Component component) {

        if (nodesToExport.isEmpty()) {
            MessageNotifyUtil.Message.info(Bundle.ExportCSV_saveNodesToCSV_empty());
            return;
        }
        
        try {   
            // Set up the file chooser with a default name and either the Export
            // folder or the last used folder.
            String fileName = getDefaultOutputFileName(nodesToExport.iterator().next().getParentNode());
            JFileChooser fileChooser = chooserHelper.getChooser();
            fileChooser.setCurrentDirectory(new File(getExportDirectory(Case.getCurrentCaseThrows())));
            fileChooser.setSelectedFile(new File(fileName));
            fileChooser.setFileFilter(new FileNameExtensionFilter("csv file", "csv"));

            int returnVal = fileChooser.showSaveDialog(component);
            if (returnVal == JFileChooser.APPROVE_OPTION) {

                // Get the file name, appending .csv if necessary
                File selectedFile = fileChooser.getSelectedFile();
                if (!selectedFile.getName().endsWith(".csv")) { // NON-NLS
                    selectedFile = new File(selectedFile.toString() + ".csv"); // NON-NLS
                }
                
                // Save the directory used for next time
                updateExportDirectory(selectedFile.getParent(), Case.getCurrentCaseThrows());
        
                if (selectedFile.exists()) {
                    logger.log(Level.SEVERE, "File {0} already exists", selectedFile.getAbsolutePath()); //NON-NLS
                    MessageNotifyUtil.Message.info(Bundle.ExportCSV_saveNodesToCSV_fileExists(selectedFile));
                    return;
                }

                CSVWriter writer = new CSVWriter(nodesToExport, selectedFile);
                writer.execute();
            }
        } catch (NoCurrentCaseException ex) {
            JOptionPane.showMessageDialog(component, Bundle.ExportCSV_saveNodesToCSV_noCurrentCase());
            logger.log(Level.INFO, "Exception while getting open case.", ex); //NON-NLS
        }
    }
    
    private static String getDefaultOutputFileName(Node parent) {
        String dateStr = String.format("%1$tY%1$tm%1$te%1$tI%1$tM%1$tS", Calendar.getInstance());
        
        if (parent != null) {
            // The first value in the property set is generally a reasonable name
            for (PropertySet set : parent.getPropertySets()) {
                for (Property<?> prop : set.getProperties()) {
                    try {
                        String parentName = prop.getValue().toString();
                        
                        // Strip off the count (if present)
                        parentName = parentName.replaceAll("\\([0-9]+\\)$", "");
                        
                        // Strip out any invalid characters
                        parentName = parentName.replaceAll("[\\\\/:*?\"<>|]", "_");
                        
                        return parentName + " " + dateStr;
                    } catch (IllegalAccessException | InvocationTargetException ex) {
                        logger.log(Level.WARNING, "Failed to get property set value as string", ex);
                    }
                }
            }
        }
        return DEFAULT_FILENAME + " " + dateStr;
    }
    
    private static String getExportDirectory(Case openCase) {
        String caseExportPath = openCase.getExportDirectory();

        if (userDefinedExportPath == null) {
            return caseExportPath;
        }

        File file = new File(userDefinedExportPath);
        if (file.exists() == false || file.isDirectory() == false) {
            return caseExportPath;
        }

        return userDefinedExportPath;
    }
    
    private static void updateExportDirectory(String exportPath, Case openCase) {
        if (exportPath.equalsIgnoreCase(openCase.getExportDirectory())) {
            userDefinedExportPath = null;
        } else {
            userDefinedExportPath = exportPath;
        }
    }

    
    private static class CSVWriter extends SwingWorker<Object, Void> {

        private static final Logger logger = Logger.getLogger(CSVWriter.class.getName());
        private ProgressHandle progress;
        
        private final Collection<? extends Node> nodesToExport;
        private final File outputFile;

        CSVWriter(Collection<? extends Node> nodesToExport, File outputFile) {
            this.nodesToExport = nodesToExport;
            this.outputFile = outputFile;
        }

        @NbBundle.Messages({"CSVWriter.progress.extracting=Exporting to CSV file",
            "CSVWriter.progress.cancelling=Cancelling"})
        @Override
        protected Object doInBackground() throws Exception {
            if (nodesToExport.isEmpty()) {
                return null;
            }

            // Set up progress bar.
            final String displayName = Bundle.CSVWriter_progress_extracting();
            progress = ProgressHandle.createHandle(displayName, new Cancellable() {
                @Override
                public boolean cancel() {
                    if (progress != null) {
                        progress.setDisplayName(Bundle.CSVWriter_progress_cancelling());
                    }
                    return ExportCSVAction.CSVWriter.this.cancel(true);
                }
            });
            progress.start();
            progress.switchToIndeterminate();

            if (this.isCancelled()) {
                return null;
            }

            Set<String> columnHeaderStrs = new HashSet<>();
            List<CsvSchema.Column> columnHeaders = new ArrayList<>();
            int remainingRowsToSample = 0;
            int columnIdx = 0;
            for (Node nd: nodesToExport) {
                // sample up to 100 rows
                if (remainingRowsToSample >= COLUMN_SAMPLING_ROW_NUM) {
                    break;
                }
                remainingRowsToSample++;
                
                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
                Iterator<?> nodeIterator = nodesToExport.iterator();
                while (nodeIterator.hasNext()) {
                    if (this.isCancelled()) {
                        return null;
                    }

                    Map<String, Object> rowMap = new HashMap<>();
                    Node node = (Node)nodeIterator.next();
                    for(PropertySet set : node.getPropertySets()) {
                        for (Property<?> prop : set.getProperties()) {
                            if (!columnsToSkip.contains(prop.getName())) {
                                rowMap.put(prop.getDisplayName(), prop.getValue());
                            }
                        }
                    }
                    seqWriter.write(rowMap);
                }
            }
            return null;
        }
        
        private String escapeQuotes(String original) {
            return original.replaceAll("\"", "\\\\\"");
        }
        
        private String listToCSV(List<String> values) {
            return "\"" + String.join("\",\"", values) + "\"\n";
        }

        @NbBundle.Messages({"CSVWriter.done.notifyMsg.error=Error exporting to CSV file",
            "# {0} - Output file",
            "CSVWriter.done.notifyMsg.success=Wrote to {0}"})
        @Override
        protected void done() {
            boolean msgDisplayed = false;
            try {
                super.get();
            } catch (InterruptedException | ExecutionException ex) {
                logger.log(Level.SEVERE, "Fatal error during file extraction", ex); //NON-NLS
                MessageNotifyUtil.Message.info(Bundle.CSVWriter_done_notifyMsg_error());
                msgDisplayed = true;
            } catch (java.util.concurrent.CancellationException ex) {
                // catch and ignore if we were cancelled
            } finally {
                progress.finish();
                if (!this.isCancelled() && !msgDisplayed) {
                    MessageNotifyUtil.Message.info(Bundle.CSVWriter_done_notifyMsg_success(outputFile));
                }
            }
        }
    }
}