CaseDbAccessManager.java

Go to the documentation of this file.

/*
 * Sleuth Kit Data Model
 *
 * Copyright 2018-2019 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.datamodel;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.sleuthkit.datamodel.SleuthkitCase.CaseDbConnection;
import org.sleuthkit.datamodel.SleuthkitCase.CaseDbTransaction;
import static org.sleuthkit.datamodel.SleuthkitCase.closeStatement;
import org.sleuthkit.datamodel.TskData.DbType;

public final class CaseDbAccessManager {
        
        public interface CaseDbAccessQueryCallback {

           void process(ResultSet resultSet);

   }


        private static final Logger logger = Logger.getLogger(CaseDbAccessManager.class.getName());

        private final SleuthkitCase tskDB;

        CaseDbAccessManager(SleuthkitCase skCase) {
                this.tskDB = skCase;
        }

        public boolean columnExists(String tableName, String columnName) throws TskCoreException {
                
                boolean doesColumnExists = false;
                CaseDbTransaction localTrans = tskDB.beginTransaction();
        try {
                        doesColumnExists = columnExists(tableName, columnName, localTrans);
                        localTrans.commit();
                        localTrans = null;
        } 
                finally {
                        if (null != localTrans) {
                                try {
                                        localTrans.rollback();
                                } catch (TskCoreException ex) {
                                        logger.log(Level.SEVERE, "Failed to rollback transaction after exception", ex);
                                }
                        }
                }
                
        return doesColumnExists;
    }
        
        public boolean columnExists(String tableName, String columnName, CaseDbTransaction transaction) throws TskCoreException {
                
                boolean columnExists = false;
        Statement statement = null;
                ResultSet resultSet = null;
        try {
                        CaseDbConnection connection = transaction.getConnection();
                        statement = connection.createStatement();
                        if (DbType.SQLITE == tskDB.getDatabaseType()) {
                                String tableInfoQuery = "PRAGMA table_info(%s)";  //NON-NLS
                                resultSet = statement.executeQuery(String.format(tableInfoQuery, tableName));
                                while (resultSet.next()) {
                                        if (resultSet.getString("name").equalsIgnoreCase(columnName)) {
                                                columnExists = true;
                                                break;
                                        }
                                }
                        }
                        else {
                                String tableInfoQueryTemplate = "SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name='%s' AND column_name='%s')";  //NON-NLS       
                                resultSet = statement.executeQuery(String.format(tableInfoQueryTemplate, tableName.toLowerCase(), columnName.toLowerCase()));
                                if (resultSet.next()) {
                                        columnExists = resultSet.getBoolean(1);
                                }
                        }
        } 
                catch (SQLException ex) {
                        throw new TskCoreException("Error checking if column  " + columnName + "exists ", ex);
                } 
                finally {
                        if (resultSet != null) {
                                try {
                                        resultSet.close();
                                } catch (SQLException ex2) {
                                        logger.log(Level.WARNING, "Failed to to close resultset after checking column", ex2);
                                }
                        }
            closeStatement(statement);
        }
        return columnExists;
    }
        
        public boolean tableExists(String tableName) throws TskCoreException {
                
                boolean doesTableExist = false;
                CaseDbTransaction localTrans = tskDB.beginTransaction();
        try {
                        doesTableExist = tableExists(tableName, localTrans);
                        localTrans.commit();
                        localTrans = null;
        } 
                finally {
                        if (null != localTrans) {
                                try {
                                        localTrans.rollback();
                                } catch (TskCoreException ex) {
                                        logger.log(Level.SEVERE, "Failed to rollback transaction after exception", ex); //NON-NLS
                                }
                        }
                }
                
        return doesTableExist;
    }
        
        public boolean tableExists(String tableName, CaseDbTransaction transaction) throws TskCoreException {
                
                boolean tableExists = false;
        Statement statement = null;
                ResultSet resultSet = null;
        try {
                        CaseDbConnection connection = transaction.getConnection();
                        statement = connection.createStatement();
                        if (DbType.SQLITE == tskDB.getDatabaseType()) {
                                resultSet = statement.executeQuery("SELECT name FROM sqlite_master WHERE type='table'");  //NON-NLS
                                while (resultSet.next()) {
                                        if (resultSet.getString("name").equalsIgnoreCase(tableName)) { //NON-NLS
                                                tableExists = true;
                                                break;
                                        }
                                }
                        }
                        else {
                                String tableInfoQueryTemplate = "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name='%s')";  //NON-NLS     
                                resultSet = statement.executeQuery(String.format(tableInfoQueryTemplate, tableName.toLowerCase()));
                                if (resultSet.next()) {
                                        tableExists = resultSet.getBoolean(1);
                                }
                        }
        } 
                catch (SQLException ex) {
                        throw new TskCoreException("Error checking if table  " + tableName + "exists ", ex);
                } finally {
                        if (resultSet != null) {
                                try {
                                        resultSet.close();
                                } catch (SQLException ex2) {
                                        logger.log(Level.WARNING, "Failed to to close resultset after checking table", ex2);
                                }
                        }
            closeStatement(statement);
        }
        return tableExists;
    }
        
        public void createTable(final String tableName, final String tableSchema) throws TskCoreException {

                validateTableName(tableName);
                validateSQL(tableSchema);

                CaseDbConnection connection = tskDB.getConnection();
                tskDB.acquireSingleUserCaseWriteLock();

                Statement statement = null;
                String createSQL = "CREATE TABLE IF NOT EXISTS " + tableName + " " + tableSchema;
                try {
                        statement = connection.createStatement();
                        statement.execute(createSQL);
                } catch (SQLException ex) {
                        throw new TskCoreException("Error creating table " + tableName, ex);
                } finally {
                        closeStatement(statement);
                        connection.close();
                        tskDB.releaseSingleUserCaseWriteLock();
                }

        }

        public void alterTable(final String tableName, final String alterSQL) throws TskCoreException {
                
                CaseDbTransaction localTrans = tskDB.beginTransaction();
                try {
                        alterTable(tableName, alterSQL, localTrans);
                        localTrans.commit();
                        localTrans = null;
                } finally {
                        if (null != localTrans) {
                                try {
                                        localTrans.rollback();
                                } catch (TskCoreException ex) {
                                        logger.log(Level.SEVERE, "Failed to rollback transaction after exception", ex);
                                }
                        }
                } 
        }
                
        public void alterTable(final String tableName, final String alterSQL, final CaseDbTransaction transaction) throws TskCoreException {

                validateTableName(tableName);
                validateSQL(alterSQL);

                CaseDbConnection connection = transaction.getConnection();
                transaction.acquireSingleUserCaseWriteLock();

                Statement statement = null;
                String sql = "ALTER TABLE " + tableName + " " + alterSQL;
                
                try {
                        statement = connection.createStatement();
                        statement.execute(sql);
                } catch (SQLException ex) {
                        // SQLite occasionally returns false for columnExists() if a table was just created with that column
                        // leading to "duplicate column name" exception.
                        // We ignore this exception
                        if (DbType.SQLITE == tskDB.getDatabaseType() &&
                                        alterSQL.toLowerCase().contains("add column") &&
                                        ex.getMessage().toLowerCase().contains("duplicate column name")) {
                                logger.log(Level.WARNING, String.format("Column being added by SQL = %s already exists in table %s", alterSQL, tableName));
                                return;
                        }
                        throw new TskCoreException(String.format("Error altering table  %s with SQL = %s", tableName, sql), ex);
                } finally {
                        closeStatement(statement);
                        // NOTE: write lock will be released by transaction
                }
        }
        
        public void createIndex(final String indexName, final String tableName, final String colsSQL) throws TskCoreException {

                validateTableName(tableName);
                validateIndexName(indexName);
                validateSQL(colsSQL);

                CaseDbConnection connection = tskDB.getConnection();
                tskDB.acquireSingleUserCaseWriteLock();

                Statement statement = null;
                String indexSQL = "CREATE INDEX IF NOT EXISTS " + indexName + " ON " + tableName + " " + colsSQL; // NON-NLS
                try {
                        statement = connection.createStatement();
                        statement.execute(indexSQL);
                } catch (SQLException ex) {
                        throw new TskCoreException("Error creating index " + tableName, ex);
                } finally {
                        closeStatement(statement);
                        connection.close();
                        tskDB.releaseSingleUserCaseWriteLock();
                }
        }

        public long insert(final String tableName, final String sql) throws TskCoreException {
                
                CaseDbTransaction localTrans = tskDB.beginTransaction();
                try {
                        long rowId = insert(tableName, sql, localTrans);
                        localTrans.commit();
                        localTrans = null;
                        return rowId;
                } finally {
                        if (null != localTrans) {
                                try {
                                        localTrans.rollback();
                                } catch (TskCoreException ex) {
                                        logger.log(Level.SEVERE, "Failed to rollback transaction after exception", ex);
                                }
                        }
                } 
                
        }
        
        public long insert(final String tableName, final String sql, final CaseDbTransaction transaction) throws TskCoreException {
                long rowId = 0;

                validateTableName(tableName);
                validateSQL(sql);

                CaseDbConnection connection = transaction.getConnection();
                transaction.acquireSingleUserCaseWriteLock();

                PreparedStatement statement = null;
                ResultSet resultSet;
                String insertSQL = "INSERT";
                if (DbType.SQLITE == tskDB.getDatabaseType()) {
                        insertSQL += " OR IGNORE";
                }
                
                insertSQL = insertSQL+ " INTO " + tableName + " " + sql; // NON-NLS
                try {
                        statement = connection.prepareStatement(insertSQL, Statement.RETURN_GENERATED_KEYS);
                        connection.executeUpdate(statement);
                        
                        resultSet = statement.getGeneratedKeys();
                        if (resultSet.next()) {
                                rowId = resultSet.getLong(1); //last_insert_rowid()
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException("Error inserting row in table " + tableName + " with sql = "+ insertSQL, ex);
                } finally {
                        closeStatement(statement);
                        // NOTE: write lock will be released by transaction
                }

                return rowId;
        }
        
        public long insertOrUpdate(final String tableName, final String sql) throws TskCoreException {
                
                CaseDbTransaction localTrans = tskDB.beginTransaction();
                try {
                        long rowId = insertOrUpdate(tableName, sql, localTrans);
                        localTrans.commit();
                        localTrans = null;
                        return rowId;
                } finally {
                        if (null != localTrans) {
                                try {
                                        localTrans.rollback();
                                } catch (TskCoreException ex) {
                                        logger.log(Level.SEVERE, "Failed to rollback transaction after exception", ex);
                                }
                        }
                } 
                
        }
        
        public long insertOrUpdate(final String tableName, final String sql, final CaseDbTransaction transaction) throws TskCoreException {
                long rowId = 0;

                validateTableName(tableName);
                validateSQL(sql);

                CaseDbConnection connection = transaction.getConnection();
                transaction.acquireSingleUserCaseWriteLock();

                PreparedStatement statement = null;
                ResultSet resultSet;
                String insertSQL = "INSERT";
                if (DbType.SQLITE == tskDB.getDatabaseType()) {
                        insertSQL += " OR REPLACE";
                }
                
                insertSQL += " INTO " + tableName + " " + sql; // NON-NLS
                try {
                        statement = connection.prepareStatement(insertSQL, Statement.RETURN_GENERATED_KEYS);
                        connection.executeUpdate(statement);

                        resultSet = statement.getGeneratedKeys();
                        resultSet.next();
                        rowId = resultSet.getLong(1); //last_insert_rowid()
                } catch (SQLException ex) {
                        throw new TskCoreException("Error inserting row in table " + tableName + " with sql = "+ insertSQL, ex);
                } finally {
                        closeStatement(statement);
                        // NOTE: write lock will be released by transaction
                }

                return rowId;
        }
        
        public void update(final String tableName, final String sql) throws TskCoreException {
                CaseDbTransaction localTrans = tskDB.beginTransaction();
                try {
                        update(tableName, sql, localTrans);
                        localTrans.commit();
                        localTrans = null;
                } finally {
                        if (null != localTrans) {
                                try {
                                        localTrans.rollback();
                                } catch (TskCoreException ex) {
                                        logger.log(Level.SEVERE, "Failed to rollback transaction after exception", ex);
                                }
                        }
                } 
        }
        

        public void update(final String tableName, final String sql, CaseDbTransaction transaction ) throws TskCoreException {
                
                validateTableName(tableName);
                validateSQL(sql);

                CaseDbConnection connection = transaction.getConnection();
                transaction.acquireSingleUserCaseWriteLock();

                Statement statement = null;
                String updateSQL = "UPDATE " + tableName + " " + sql; // NON-NLS

                try {
                        statement = connection.createStatement();
                        statement.executeUpdate(updateSQL);
                } catch (SQLException ex) {
                        throw new TskCoreException("Error Updating table " + tableName, ex);
                } finally {
                        closeStatement(statement);
                        // NOTE: write lock will be released by transaction
                }
        }
        
        public void select(final String sql, final CaseDbAccessQueryCallback queryCallback) throws TskCoreException {

                if (queryCallback == null) {
            throw new TskCoreException("Callback is null");
        }
                
                validateSQL(sql);
                
                CaseDbConnection connection = tskDB.getConnection();
                tskDB.acquireSingleUserCaseReadLock();

                Statement statement = null;
                ResultSet resultSet;
                String selectSQL = "SELECT " +  sql; // NON-NLS
                try {
                        statement = connection.createStatement();
                        resultSet = statement.executeQuery(selectSQL);
                        queryCallback.process(resultSet);
                } catch (SQLException ex) {
                        throw new TskCoreException("Error running SELECT query.", ex);
                } finally {
                        closeStatement(statement);
                        connection.close();
                        tskDB.releaseSingleUserCaseReadLock();
                }
        }
        
        public void delete(final String tableName, final String sql ) throws TskCoreException {
                validateTableName(tableName);
                validateSQL(sql);

                CaseDbConnection connection = tskDB.getConnection();
                tskDB.acquireSingleUserCaseWriteLock();

                Statement statement = null;
                String deleteSQL = "DELETE FROM " + tableName + " " + sql; // NON-NLS
                try {
                        statement = connection.createStatement();
                        statement.executeUpdate(deleteSQL);
                } catch (SQLException ex) {
                        throw new TskCoreException("Error deleting row from table " + tableName, ex);
                } finally {
                        closeStatement(statement);
                        connection.close();
                        tskDB.releaseSingleUserCaseWriteLock();
                }
        }

        private void validateTableName(String tableName) throws TskCoreException {
                
                if (SleuthkitCase.getCoreTableNames().contains(tableName.toLowerCase())) {
                        throw new TskCoreException("Attempt to modify a core TSK table " + tableName);
                }
                if (tableName.toLowerCase().startsWith("tsk_")) {
                        throw new TskCoreException("Modifying tables with tsk_ prefix is not allowed. ");
                }
        }

        private void validateIndexName(String indexName) throws TskCoreException {
                
                if (indexName.isEmpty()) {
                        throw new TskCoreException("Invalid index name " + indexName);  
                }
                
                if (SleuthkitCase.getCoreIndexNames().contains(indexName.toLowerCase())) {
                        throw new TskCoreException("Attempt to modify a core TSK index " + indexName);  
                }
        }
        
        private void validateSQL(String sql) throws TskCoreException {
                /*
                 * TODO (JIRA-5950): Need SQL injection defense in CaseDbAccessManager 
                 */
        }

}