IngestTasksScheduler.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 * 
 * Copyright 2012-2015 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.ingest;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.logging.Level;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.sleuthkit.autopsy.coreutils.Logger;
import org.sleuthkit.datamodel.AbstractFile;
import org.sleuthkit.datamodel.Content;
import org.sleuthkit.datamodel.FileSystem;
import org.sleuthkit.datamodel.TskCoreException;
import org.sleuthkit.datamodel.TskData;

final class IngestTasksScheduler {

    private static final Logger logger = Logger.getLogger(IngestTasksScheduler.class.getName());
    private static final int FAT_NTFS_FLAGS = TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_FAT12.getValue() | TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_FAT16.getValue() | TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_FAT32.getValue() | TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_NTFS.getValue();
    private static IngestTasksScheduler instance;

    private final LinkedBlockingQueue<DataSourceIngestTask> pendingDataSourceTasks;
    private final DataSourceIngestTaskQueue dataSourceTasksDispenser;

    private final TreeSet<FileIngestTask> rootDirectoryTasks;
    private final List<FileIngestTask> directoryTasks;
    private final BlockingDeque<FileIngestTask> pendingFileTasks;
    private final FileIngestTaskQueue fileTasksDispenser;

    private final Set<IngestTask> tasksInProgress;

    synchronized static IngestTasksScheduler getInstance() {
        if (IngestTasksScheduler.instance == null) {
            IngestTasksScheduler.instance = new IngestTasksScheduler();
        }
        return IngestTasksScheduler.instance;
    }

    private IngestTasksScheduler() {
        this.pendingDataSourceTasks = new LinkedBlockingQueue<>();
        this.dataSourceTasksDispenser = new DataSourceIngestTaskQueue();
        this.rootDirectoryTasks = new TreeSet<>(new RootDirectoryTaskComparator());
        this.directoryTasks = new ArrayList<>();
        this.pendingFileTasks = new LinkedBlockingDeque<>();
        this.fileTasksDispenser = new FileIngestTaskQueue();
        this.tasksInProgress = new HashSet<>();
    }

    IngestTaskQueue getDataSourceIngestTaskQueue() {
        return this.dataSourceTasksDispenser;
    }

    IngestTaskQueue getFileIngestTaskQueue() {
        return this.fileTasksDispenser;
    }

    synchronized void scheduleIngestTasks(DataSourceIngestJob job) {
        if (!job.isCancelled()) {
            // Scheduling of both a data source ingest task and file ingest tasks 
            // for a job must be an atomic operation. Otherwise, the data source 
            // task might be completed before the file tasks are scheduled, 
            // resulting in a potential false positive when another thread checks 
            // whether or not all the tasks for the job are completed.
            this.scheduleDataSourceIngestTask(job);
            this.scheduleFileIngestTasks(job);
        }
    }

    synchronized void scheduleDataSourceIngestTask(DataSourceIngestJob job) {
        if (!job.isCancelled()) {
            DataSourceIngestTask task = new DataSourceIngestTask(job);
            this.tasksInProgress.add(task);
            try {
                this.pendingDataSourceTasks.put(task);
            } catch (InterruptedException ex) {
                this.tasksInProgress.remove(task);
                Thread.currentThread().interrupt();
            }
        }
    }

    synchronized void scheduleFileIngestTasks(DataSourceIngestJob job) {
        if (!job.isCancelled()) {
            // Get the top level files for the data source associated with this job
            // and add them to the root directories priority queue.  
            List<AbstractFile> topLevelFiles = getTopLevelFiles(job.getDataSource());
            for (AbstractFile firstLevelFile : topLevelFiles) {
                FileIngestTask task = new FileIngestTask(job, firstLevelFile);
                if (IngestTasksScheduler.shouldEnqueueFileTask(task)) {
                    this.tasksInProgress.add(task);
                    this.rootDirectoryTasks.add(task);
                }
            }
            shuffleFileTaskQueues();
        }
    }

    synchronized void scheduleFileIngestTask(DataSourceIngestJob job, AbstractFile file) {
        if (!job.isCancelled()) {
            FileIngestTask task = new FileIngestTask(job, file);
            if (IngestTasksScheduler.shouldEnqueueFileTask(task)) {
                this.tasksInProgress.add(task);
                addToPendingFileTasksQueue(task);
            }
        }
    }

    synchronized void notifyTaskCompleted(IngestTask task) {
        tasksInProgress.remove(task);
    }

    synchronized boolean tasksForJobAreCompleted(DataSourceIngestJob job) {
        for (IngestTask task : tasksInProgress) {
            if (task.getIngestJob().getId() == job.getId()) {
                return false;
            }
        }
        return true;
    }

    synchronized void cancelPendingTasksForIngestJob(DataSourceIngestJob job) {
        long jobId = job.getId();
        this.removeTasksForJob(this.rootDirectoryTasks, jobId);
        this.removeTasksForJob(this.directoryTasks, jobId);
        this.shuffleFileTaskQueues();
    }

    private static List<AbstractFile> getTopLevelFiles(Content dataSource) {
        List<AbstractFile> topLevelFiles = new ArrayList<>();
        Collection<AbstractFile> rootObjects = dataSource.accept(new GetRootDirectoryVisitor());
        if (rootObjects.isEmpty() && dataSource instanceof AbstractFile) {
            // The data source is itself a file to be processed.
            topLevelFiles.add((AbstractFile) dataSource);
        } else {
            for (AbstractFile root : rootObjects) {
                List<Content> children;
                try {
                    children = root.getChildren();
                    if (children.isEmpty()) {
                        // Add the root object itself, it could be an unallocated
                        // space file, or a child of a volume or an image.
                        topLevelFiles.add(root);
                    } else {
                        // The root object is a file system root directory, get
                        // the files within it.
                        for (Content child : children) {
                            if (child instanceof AbstractFile) {
                                topLevelFiles.add((AbstractFile) child);
                            }
                        }
                    }
                } catch (TskCoreException ex) {
                    logger.log(Level.WARNING, "Could not get children of root to enqueue: " + root.getId() + ": " + root.getName(), ex); //NON-NLS
                }
            }
        }
        return topLevelFiles;
    }

    synchronized private void shuffleFileTaskQueues() {
        // This is synchronized because it is called both by synchronized 
        // methods of this ingest scheduler and an unsynchronized method of its
        // file tasks "dispenser".
        while (true) {
            // Loop until either the pending file tasks queue is NOT empty
            // or the upstream queues that feed into it ARE empty.
            if (!this.pendingFileTasks.isEmpty()) {
                // There are file tasks ready to be consumed, exit.
                return;
            }
            if (this.directoryTasks.isEmpty()) {
                if (this.rootDirectoryTasks.isEmpty()) {
                    // There are no root directory tasks to move into the
                    // directory queue, exit.
                    return;
                } else {
                    // Move the next root directory task into the 
                    // directories queue. Note that the task was already 
                    // added to the tasks in progress list when the task was
                    // created in scheduleFileIngestTasks().
                    this.directoryTasks.add(this.rootDirectoryTasks.pollFirst());
                }
            }

            // Try to add the most recently added directory from the 
            // directory tasks queue to the pending file tasks queue. 
            FileIngestTask directoryTask = this.directoryTasks.remove(this.directoryTasks.size() - 1);
            if (shouldEnqueueFileTask(directoryTask)) {
                addToPendingFileTasksQueue(directoryTask);
            } else {
                this.tasksInProgress.remove(directoryTask);
            }

            // If the directory contains subdirectories or files, try to 
            // enqueue tasks for them as well. 
            final AbstractFile directory = directoryTask.getFile();
            try {
                for (Content child : directory.getChildren()) {
                    if (child instanceof AbstractFile) {
                        AbstractFile file = (AbstractFile) child;
                        FileIngestTask childTask = new FileIngestTask(directoryTask.getIngestJob(), file);
                        if (file.hasChildren()) {
                            // Found a subdirectory, put the task in the 
                            // pending directory tasks queue. Note the 
                            // addition of the task to the tasks in progress
                            // list. This is necessary because this is the
                            // first appearance of this task in the queues.
                            this.tasksInProgress.add(childTask);
                            this.directoryTasks.add(childTask);
                        } else if (shouldEnqueueFileTask(childTask)) {
                            // Found a file, put the task directly into the
                            // pending file tasks queue. 
                            this.tasksInProgress.add(childTask);
                            addToPendingFileTasksQueue(childTask);
                        }
                    }
                }
            } catch (TskCoreException ex) {
                String errorMessage = String.format("An error occurred getting the children of %s", directory.getName()); //NON-NLS
                logger.log(Level.SEVERE, errorMessage, ex);
            }
        }
    }

    private static boolean shouldEnqueueFileTask(final FileIngestTask task) {
        final AbstractFile file = task.getFile();

        // Skip the task if the file is an unallocated space file and the
        // process unallocated space flag is not set for this job.
        if (!task.getIngestJob().shouldProcessUnallocatedSpace()
                && file.getType().equals(TskData.TSK_DB_FILES_TYPE_ENUM.UNALLOC_BLOCKS)) {
            return false;
        }

        // Skip the task if the file is actually the pseudo-file for the parent
        // or current directory.
        String fileName = file.getName();
        if (fileName.equals(".") || fileName.equals("..")) {
            return false;
        }

        // Skip the task if the file is one of a select group of special, large
        // NTFS or FAT file system files.
        if (file instanceof org.sleuthkit.datamodel.File) {
            final org.sleuthkit.datamodel.File f = (org.sleuthkit.datamodel.File) file;

            // Get the type of the file system, if any, that owns the file.
            TskData.TSK_FS_TYPE_ENUM fsType = TskData.TSK_FS_TYPE_ENUM.TSK_FS_TYPE_UNSUPP;
            try {
                FileSystem fs = f.getFileSystem();
                if (fs != null) {
                    fsType = fs.getFsType();
                }
            } catch (TskCoreException ex) {
                logger.log(Level.SEVERE, "Error querying file system for " + f, ex); //NON-NLS
            }

            // If the file system is not NTFS or FAT, don't skip the file.
            if ((fsType.getValue() & FAT_NTFS_FLAGS) == 0) {
                return true;
            }

            // Find out whether the file is in a root directory. 
            boolean isInRootDir = false;
            try {
                AbstractFile parent = f.getParentDirectory();
                isInRootDir = parent.isRoot();
            } catch (TskCoreException ex) {
                logger.log(Level.WARNING, "Error querying parent directory for" + f.getName(), ex); //NON-NLS
            }

            // If the file is in the root directory of an NTFS or FAT file 
            // system, check its meta-address and check its name for the '$'
            // character and a ':' character (not a default attribute).
            if (isInRootDir && f.getMetaAddr() < 32) {
                String name = f.getName();
                if (name.length() > 0 && name.charAt(0) == '$' && name.contains(":")) {
                    return false;
                }
            }
        }

        return true;
    }

    synchronized private void addToPendingFileTasksQueue(FileIngestTask task) {
        try {
            this.pendingFileTasks.putFirst(task);
        } catch (InterruptedException ex) {
            this.tasksInProgress.remove(task);
            Thread.currentThread().interrupt();
        }
    }

    synchronized private void removeTasksForJob(Collection<? extends IngestTask> taskQueue, long jobId) {
        Iterator<? extends IngestTask> iterator = taskQueue.iterator();
        while (iterator.hasNext()) {
            IngestTask task = iterator.next();
            if (task.getIngestJob().getId() == jobId) {
                this.tasksInProgress.remove(task);
                iterator.remove();
            }
        }
    }

    private static int countTasksForJob(Collection<? extends IngestTask> queue, long jobId) {
        Iterator<? extends IngestTask> iterator = queue.iterator();
        int count = 0;
        while (iterator.hasNext()) {
            IngestTask task = (IngestTask) iterator.next();
            if (task.getIngestJob().getId() == jobId) {
                count++;
            }
        }
        return count;
    }

    synchronized IngestJobTasksSnapshot getTasksSnapshotForJob(long jobId) {
        return new IngestJobTasksSnapshot(jobId);
    }

    private static class RootDirectoryTaskComparator implements Comparator<FileIngestTask> {

        @Override
        public int compare(FileIngestTask q1, FileIngestTask q2) {
            AbstractFilePriority.Priority p1 = AbstractFilePriority.getPriority(q1.getFile());
            AbstractFilePriority.Priority p2 = AbstractFilePriority.getPriority(q2.getFile());
            if (p1 == p2) {
                return (int) (q2.getFile().getId() - q1.getFile().getId());
            } else {
                return p2.ordinal() - p1.ordinal();
            }
        }

        private static class AbstractFilePriority {

            enum Priority {

                LAST, LOW, MEDIUM, HIGH
            }

            static final List<Pattern> LAST_PRI_PATHS = new ArrayList<>();

            static final List<Pattern> LOW_PRI_PATHS = new ArrayList<>();

            static final List<Pattern> MEDIUM_PRI_PATHS = new ArrayList<>();

            static final List<Pattern> HIGH_PRI_PATHS = new ArrayList<>();
            /*
             * prioritize root directory folders based on the assumption that we
             * are looking for user content. Other types of investigations may
             * want different priorities.
             */

            static /*
             * prioritize root directory folders based on the assumption that we
             * are looking for user content. Other types of investigations may
             * want different priorities.
             */ {
                // these files have no structure, so they go last
                //unalloc files are handled as virtual files in getPriority()
                //LAST_PRI_PATHS.schedule(Pattern.compile("^\\$Unalloc", Pattern.CASE_INSENSITIVE));
                //LAST_PRI_PATHS.schedule(Pattern.compile("^\\Unalloc", Pattern.CASE_INSENSITIVE));
                LAST_PRI_PATHS.add(Pattern.compile("^pagefile", Pattern.CASE_INSENSITIVE));
                LAST_PRI_PATHS.add(Pattern.compile("^hiberfil", Pattern.CASE_INSENSITIVE));
                // orphan files are often corrupt and windows does not typically have
                // user content, so put them towards the bottom
                LOW_PRI_PATHS.add(Pattern.compile("^\\$OrphanFiles", Pattern.CASE_INSENSITIVE));
                LOW_PRI_PATHS.add(Pattern.compile("^Windows", Pattern.CASE_INSENSITIVE));
                // all other files go into the medium category too
                MEDIUM_PRI_PATHS.add(Pattern.compile("^Program Files", Pattern.CASE_INSENSITIVE));
                // user content is top priority
                HIGH_PRI_PATHS.add(Pattern.compile("^Users", Pattern.CASE_INSENSITIVE));
                HIGH_PRI_PATHS.add(Pattern.compile("^Documents and Settings", Pattern.CASE_INSENSITIVE));
                HIGH_PRI_PATHS.add(Pattern.compile("^home", Pattern.CASE_INSENSITIVE));
                HIGH_PRI_PATHS.add(Pattern.compile("^ProgramData", Pattern.CASE_INSENSITIVE));
            }

            static AbstractFilePriority.Priority getPriority(final AbstractFile abstractFile) {
                if (!abstractFile.getType().equals(TskData.TSK_DB_FILES_TYPE_ENUM.FS)) {
                    //quickly filter out unstructured content
                    //non-fs virtual files and dirs, such as representing unalloc space
                    return AbstractFilePriority.Priority.LAST;
                }
                //determine the fs files priority by name
                final String path = abstractFile.getName();
                if (path == null) {
                    return AbstractFilePriority.Priority.MEDIUM;
                }
                for (Pattern p : HIGH_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.HIGH;
                    }
                }
                for (Pattern p : MEDIUM_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.MEDIUM;
                    }
                }
                for (Pattern p : LOW_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.LOW;
                    }
                }
                for (Pattern p : LAST_PRI_PATHS) {
                    Matcher m = p.matcher(path);
                    if (m.find()) {
                        return AbstractFilePriority.Priority.LAST;
                    }
                }
                //default is medium
                return AbstractFilePriority.Priority.MEDIUM;
            }
        }
    }

    private final class DataSourceIngestTaskQueue implements IngestTaskQueue {

        @Override
        public IngestTask getNextTask() throws InterruptedException {
            return IngestTasksScheduler.this.pendingDataSourceTasks.take();
        }
    }

    private final class FileIngestTaskQueue implements IngestTaskQueue {

        @Override
        public IngestTask getNextTask() throws InterruptedException {
            FileIngestTask task = IngestTasksScheduler.this.pendingFileTasks.takeFirst();
            shuffleFileTaskQueues();
            return task;
        }

    }

    class IngestJobTasksSnapshot {

        private final long jobId;
        private final long rootQueueSize;
        private final long dirQueueSize;
        private final long fileQueueSize;
        private final long dsQueueSize;
        private final long runningListSize;

        IngestJobTasksSnapshot(long jobId) {
            this.jobId = jobId;
            this.rootQueueSize = countTasksForJob(IngestTasksScheduler.this.rootDirectoryTasks, jobId);
            this.dirQueueSize = countTasksForJob(IngestTasksScheduler.this.directoryTasks, jobId);
            this.fileQueueSize = countTasksForJob(IngestTasksScheduler.this.pendingFileTasks, jobId);
            this.dsQueueSize = countTasksForJob(IngestTasksScheduler.this.pendingDataSourceTasks, jobId);
            this.runningListSize = countTasksForJob(IngestTasksScheduler.this.tasksInProgress, jobId);
        }

        long getJobId() {
            return jobId;
        }

        long getRootQueueSize() {
            return rootQueueSize;
        }

        long getDirectoryTasksQueueSize() {
            return dirQueueSize;
        }

        long getFileQueueSize() {
            return fileQueueSize;
        }

        long getDsQueueSize() {
            return dsQueueSize;
        }

        long getRunningListSize() {
            return runningListSize;
        }

    }

}