SQLiteTableReader.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2018-2018 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.coreutils;

import java.io.File;
import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.function.BooleanSupplier;
import java.util.function.Consumer;
import java.util.logging.Level;
import org.sleuthkit.autopsy.casemodule.Case;
import org.sleuthkit.autopsy.casemodule.NoCurrentCaseException;
import org.sleuthkit.autopsy.casemodule.services.FileManager;
import org.sleuthkit.autopsy.casemodule.services.Services;
import org.sleuthkit.autopsy.datamodel.ContentUtils;
import org.sleuthkit.datamodel.AbstractFile;
import org.sleuthkit.datamodel.SleuthkitCase;
import org.sleuthkit.datamodel.TskCoreException;

public class SQLiteTableReader implements AutoCloseable {

    public static class Builder {

        private final AbstractFile file;

        private Consumer<String> forAllColumnNamesConsumer;
        private Consumer<String> forAllStringValuesConsumer;
        private Consumer<Long> forAllLongValuesConsumer;
        private Consumer<Integer> forAllIntegerValuesConsumer;
        private Consumer<Double> forAllFloatValuesConsumer;
        private Consumer<byte[]> forAllBlobValuesConsumer;
        private Consumer<Object> forAllTableValuesConsumer;
        
        static <T> Consumer<T> doNothing() {
            return NOOP -> {};
        }

        public Builder(AbstractFile file) {
            this.file = file;

            this.forAllColumnNamesConsumer = Builder.doNothing();
            this.forAllStringValuesConsumer = Builder.doNothing();
            this.forAllLongValuesConsumer = Builder.doNothing();
            this.forAllIntegerValuesConsumer = Builder.doNothing();
            this.forAllFloatValuesConsumer = Builder.doNothing();
            this.forAllBlobValuesConsumer = Builder.doNothing();
            this.forAllTableValuesConsumer = Builder.doNothing();
        }

        public Builder forAllColumnNames(Consumer<String> action) {
            this.forAllColumnNamesConsumer = action;
            return this;
        }

        public Builder forAllStringValues(Consumer<String> action) {
            this.forAllStringValuesConsumer = action;
            return this;
        }

        public Builder forAllIntegerValues(Consumer<Integer> action) {
            this.forAllIntegerValuesConsumer = action;
            return this;
        }

        public Builder forAllFloatValues(Consumer<Double> action) {
            this.forAllFloatValuesConsumer = action;
            return this;
        }

        public Builder forAllLongValues(Consumer<Long> action) {
            this.forAllLongValuesConsumer = action;
            return this;
        }

        public Builder forAllBlobValues(Consumer<byte[]> action) {
            this.forAllBlobValuesConsumer = action;
            return this;
        }

        public Builder forAllTableValues(Consumer<Object> action) {
            this.forAllTableValuesConsumer = action;
            return this;
        }

        public SQLiteTableReader build() {
            return new SQLiteTableReader(this);
        }
    }

    private final AbstractFile file;
    private final Builder builder;
  
    private static final String SELECT_ALL_QUERY = "SELECT * FROM \"%s\"";
    private static final Logger logger = Logger.getLogger(SQLiteTableReader.class.getName());

    private Connection conn;
    private PreparedStatement statement;
    private ResultSet queryResults;
    private ResultSetMetaData currentMetadata;

    //Iteration state
    private int currRowColumnIndex;
    private int columnNameIndex;
    private int totalColumnCount;
    private boolean unfinishedRow;
    private boolean liveResultSet;
    private String prevTableName;

    private SQLiteTableReader(Builder builder) {
        this.builder = builder;
        this.file = builder.file;
    }

    public List<String> getTableNames() throws SQLiteTableReaderException {
        ensureOpen();
        try (ResultSet tableNameResult = conn.createStatement()
                .executeQuery("SELECT name FROM sqlite_master "
                        + " WHERE type= 'table' ")) {
            List<String> tableNames = new ArrayList<>();
            while (tableNameResult.next()) {
                tableNames.add(tableNameResult.getString("name")); //NON-NLS
            }
            return tableNames;
        } catch (SQLException ex) {
            throw new SQLiteTableReaderException(ex);
        }
    }

    public int getRowCount(String tableName) throws SQLiteTableReaderException {
        ensureOpen();
        try (ResultSet countResult = conn.createStatement()
                .executeQuery("SELECT count (*) as count FROM "
                        + "\"" + tableName + "\"")) {
            return countResult.getInt("count");
        } catch (SQLException ex) {
            throw new SQLiteTableReaderException(ex);
        }
    }

    public int getColumnCount(String tableName) throws SQLiteTableReaderException {
        ensureOpen();
        try (ResultSet columnCount = conn.createStatement()
                .executeQuery(String.format(SELECT_ALL_QUERY, tableName))) {
            return columnCount.getMetaData().getColumnCount();
        } catch (SQLException ex) {
            throw new SQLiteTableReaderException(ex);
        }
    }

    public void read(String tableName) throws SQLiteTableReaderException {
        readHelper(String.format(SELECT_ALL_QUERY, tableName), () -> false);
    }

    public void read(String tableName, int limit, int offset) throws SQLiteTableReaderException {
        readHelper(String.format(SELECT_ALL_QUERY, tableName) + " LIMIT " + limit
                + " OFFSET " + offset, () -> false);
    }

    public void read(String tableName, BooleanSupplier condition) throws SQLiteTableReaderException {
        if (Objects.isNull(prevTableName) || !prevTableName.equals(tableName)) {
            prevTableName = tableName;
            closeTableResources();
        }
        readHelper(String.format(SELECT_ALL_QUERY, tableName), condition);
    }

    private void readHelper(String query, BooleanSupplier condition) throws SQLiteTableReaderException {
        try {
            if (!liveResultSet) {
                openTableResources(query);
                columnNameIndex = 0;
            }

            //Process column names before reading the database table values
            while (columnNameIndex < totalColumnCount) {
                if (condition.getAsBoolean()) {
                    return;
                }
                builder.forAllColumnNamesConsumer.accept(currentMetadata
                        .getColumnName(++columnNameIndex));
            }

            while (unfinishedRow || queryResults.next()) {
                while (currRowColumnIndex < totalColumnCount) {
                    if (condition.getAsBoolean()) {
                        unfinishedRow = true;
                        return;
                    }

                    Object item = queryResults.getObject(++currRowColumnIndex);
                    if (item instanceof String) {
                        builder.forAllStringValuesConsumer.accept((String) item);
                    } else if (item instanceof Integer) {
                        builder.forAllIntegerValuesConsumer.accept((Integer) item);
                    } else if (item instanceof Double) {
                        builder.forAllFloatValuesConsumer.accept((Double) item);
                    } else if (item instanceof Long) {
                        builder.forAllLongValuesConsumer.accept((Long) item);
                    } else if (item instanceof byte[]) {
                        builder.forAllBlobValuesConsumer.accept((byte[]) item);
                    }

                    builder.forAllTableValuesConsumer.accept(item);
                }
                unfinishedRow = false;
                //Wrap column index back around if we've reached the end of the row
                currRowColumnIndex = currRowColumnIndex % totalColumnCount;
            }
            closeTableResources();
        } catch (SQLException ex) {
            closeTableResources();
            throw new SQLiteTableReaderException(ex);
        }
    }

    private void ensureOpen() throws SQLiteTableReaderException {
        if (Objects.isNull(conn)) {
            try {
                Class.forName("org.sqlite.JDBC"); //NON-NLS  
                String localDiskPath = copyFileToTempDirectory(file, file.getId());

                //Find and copy both WAL and SHM meta files 
                findAndCopySQLiteMetaFile(file, file.getName() + "-wal");
                findAndCopySQLiteMetaFile(file, file.getName() + "-shm");
                conn = DriverManager.getConnection("jdbc:sqlite:" + localDiskPath);
            } catch (NoCurrentCaseException | TskCoreException | IOException
                    | ClassNotFoundException | SQLException ex) {
                throw new SQLiteTableReaderException(ex);
            }
        }
    }

    private void findAndCopySQLiteMetaFile(AbstractFile sqliteFile,
            String metaFileName) throws NoCurrentCaseException, TskCoreException, IOException {

        Case openCase = Case.getCurrentCaseThrows();
        SleuthkitCase sleuthkitCase = openCase.getSleuthkitCase();
        Services services = new Services(sleuthkitCase);
        FileManager fileManager = services.getFileManager();

        List<AbstractFile> metaFiles = fileManager.findFiles(
                sqliteFile.getDataSource(), metaFileName,
                sqliteFile.getParent().getName());

        if (metaFiles != null) {
            for (AbstractFile metaFile : metaFiles) {
                copyFileToTempDirectory(metaFile, sqliteFile.getId());
            }
        }
    }

    private String copyFileToTempDirectory(AbstractFile file, long fileId)
            throws IOException, NoCurrentCaseException {

        String localDiskPath = Case.getCurrentCaseThrows().getTempDirectory()
                + File.separator + fileId + file.getName();
        File localDatabaseFile = new File(localDiskPath);
        if (!localDatabaseFile.exists()) {
            ContentUtils.writeToFile(file, localDatabaseFile);
        }
        return localDiskPath;
    }

    private void openTableResources(String query) throws SQLiteTableReaderException {
        try {
            ensureOpen();
            statement = conn.prepareStatement(query);
            queryResults = statement.executeQuery();
            currentMetadata = queryResults.getMetaData();
            totalColumnCount = currentMetadata.getColumnCount();
            liveResultSet = true;
        } catch (SQLException ex) {
            throw new SQLiteTableReaderException(ex);
        }
    }

    private void closeTableResources() {
        try {
            if (Objects.nonNull(statement)) {
                statement.close();
            }
            if (Objects.nonNull(queryResults)) {
                queryResults.close();
            }
            liveResultSet = false;
        } catch (SQLException ex) {
            logger.log(Level.SEVERE, "Failed to close table resources", ex);
        }
    }

    @Override
    public void close() throws SQLiteTableReaderException {
        try {
            if (Objects.nonNull(conn)) {
                conn.close();
            }
        } catch (SQLException ex) {
            throw new SQLiteTableReaderException(ex);
        }
    }

    public boolean isFinished() {
        return !liveResultSet;
    }

    @Override
    protected void finalize() throws Throwable {
        try {
            close();
        } catch (SQLiteTableReaderException ex) {
            logger.log(Level.SEVERE, "Failed to close reader in finalizer", ex);
        }
        super.finalize();
    }
}