ExcelExport.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 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 file 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.datasourcesummary.uiutils;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.CellStyle;
import org.apache.poi.ss.usermodel.Font;
import org.apache.poi.ss.usermodel.HorizontalAlignment;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.openide.util.NbBundle.Messages;
import org.sleuthkit.autopsy.datasourcesummary.uiutils.CellModel.HorizontalAlign;

public class ExcelExport {

    public static class ExcelExportException extends Exception {

        public ExcelExportException(String string) {
            super(string);
        }

        public ExcelExportException(String string, Throwable thrwbl) {
            super(string, thrwbl);
        }
    }

    static class CellStyleKey {

        private final String formatString;
        private final CellStyle cellStyle;
        private final HorizontalAlign alignment;

        CellStyleKey(String formatString, CellStyle cellStyle, HorizontalAlign alignment) {
            this.formatString = formatString;
            this.cellStyle = cellStyle;
            this.alignment = alignment;
        }

        String getFormatString() {
            return formatString;
        }

        CellStyle getCellStyle() {
            return cellStyle;
        }

        HorizontalAlign getAlignment() {
            return alignment;
        }

        @Override
        public int hashCode() {
            int hash = 7;
            hash = 29 * hash + Objects.hashCode(this.formatString);
            hash = 29 * hash + Objects.hashCode(this.cellStyle);
            hash = 29 * hash + Objects.hashCode(this.alignment);
            return hash;
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            if (obj == null) {
                return false;
            }
            if (getClass() != obj.getClass()) {
                return false;
            }
            final CellStyleKey other = (CellStyleKey) obj;
            if (!Objects.equals(this.formatString, other.formatString)) {
                return false;
            }
            if (!Objects.equals(this.cellStyle, other.cellStyle)) {
                return false;
            }
            if (this.alignment != other.alignment) {
                return false;
            }
            return true;
        }
    }

    public static class WorksheetEnv {

        private final CellStyle headerStyle;
        private final Workbook parentWorkbook;
        private final CellStyle defaultStyle;

        // maps a data format string / original cell style combination to a created cell style
        private final Map<CellStyleKey, CellStyle> cellStyleCache = new HashMap<>();

        WorksheetEnv(CellStyle headerStyle, CellStyle defaultStyle, Workbook parentWorkbook) {
            this.headerStyle = headerStyle;
            this.defaultStyle = defaultStyle;
            this.parentWorkbook = parentWorkbook;
        }

        public CellStyle getCellStyle(CellStyleKey cellStyleKey) {
            return cellStyleCache.computeIfAbsent(cellStyleKey, (pair) -> {
                CellStyle computed = this.parentWorkbook.createCellStyle();
                computed.cloneStyleFrom(cellStyleKey.getCellStyle() == null ? defaultStyle : cellStyleKey.getCellStyle());

                if (cellStyleKey.getAlignment() != null) {
                    computed.setAlignment(cellStyleKey.getAlignment().getPoiAlignment());
                }

                if (cellStyleKey.getFormatString() != null) {
                    computed.setDataFormat(this.parentWorkbook.getCreationHelper().createDataFormat().getFormat(cellStyleKey.getFormatString()));
                }
                return computed;
            });
        }

        public CellStyle getHeaderStyle() {
            return headerStyle;
        }

        public CellStyle getDefaultCellStyle() {
            return defaultStyle;
        }

        public Workbook getParentWorkbook() {
            return parentWorkbook;
        }
    }

    public static interface ExcelSheetExport {

        String getSheetName();

        void renderSheet(Sheet sheet, WorksheetEnv env) throws ExcelExportException;
    }

    private static ExcelExport instance = null;

    public static ExcelExport getInstance() {
        if (instance == null) {
            instance = new ExcelExport();
        }

        return instance;
    }

    private ExcelExport() {

    }

    @Messages({
        "# {0} - sheetNumber",
        "ExcelExport_writeExcel_noSheetName=Sheet {0}"
    })
    public void writeExcel(List<ExcelSheetExport> exports, File path) throws IOException, ExcelExportException {
        // Create a Workbook
        Workbook workbook = new XSSFWorkbook(); // new HSSFWorkbook() for generating `.xls` file

        // Create a Font for styling header cells
        Font headerFont = workbook.createFont();
        headerFont.setBold(true);
        //headerFont.setFontHeightInPoints((short) 14);

        // Create a CellStyle with the font
        HorizontalAlignment alignment = HorizontalAlignment.LEFT;
        CellStyle headerCellStyle = workbook.createCellStyle();
        headerCellStyle.setFont(headerFont);
        headerCellStyle.setAlignment(alignment);

        CellStyle defaultCellStyle = workbook.createCellStyle();
        defaultCellStyle.setAlignment(alignment);

        WorksheetEnv env = new WorksheetEnv(headerCellStyle, defaultCellStyle, workbook);

        if (exports != null) {
            for (int i = 0; i < exports.size(); i++) {
                ExcelSheetExport export = exports.get(i);
                if (export == null) {
                    continue;
                }

                String sheetName = export.getSheetName();
                if (sheetName == null) {
                    sheetName = Bundle.ExcelExport_writeExcel_noSheetName(i + 1);
                }

                Sheet sheet = workbook.createSheet(sheetName);
                export.renderSheet(sheet, env);
            }
        }

        // Write the output to a file
        FileOutputStream fileOut = new FileOutputStream(path);
        workbook.write(fileOut);
        fileOut.close();

        // Closing the workbook
        workbook.close();
    }

    static Cell createCell(WorksheetEnv env, Row row, int colNum, ExcelCellModel cellModel, Optional<CellStyle> cellStyle) {
        CellStyle cellStyleToUse = cellStyle.orElse(env.getDefaultCellStyle());

        if (cellModel.getExcelFormatString() != null || cellModel.getHorizontalAlignment() != null) {
            cellStyleToUse = env.getCellStyle(new CellStyleKey(cellModel.getExcelFormatString(), cellStyleToUse, cellModel.getHorizontalAlignment()));
        }

        Object cellData = cellModel.getData();
        Cell cell = row.createCell(colNum);
        if (cellData instanceof Calendar) {
            cell.setCellValue((Calendar) cellData);
        } else if (cellData instanceof Date) {
            cell.setCellValue((Date) cellData);
        } else if (cellData instanceof Double) {
            cell.setCellValue((Double) cellData);
        } else if (cellData instanceof String) {
            cell.setCellValue((String) cellData);
        } else if (cellData instanceof Short) {
            cell.setCellValue((Short) cellData);
        } else if (cellData instanceof Integer) {
            cell.setCellValue((Integer) cellData);
        } else if (cellData instanceof Long) {
            cell.setCellValue((Long) cellData);
        } else if (cellData instanceof Float) {
            cell.setCellValue((Float) cellData);
        } else {
            cell.setCellValue(cellModel.getText());
        }
        cell.setCellStyle(cellStyleToUse);
        return cell;
    }
}