EmailMessageThreader.java

Go to the documentation of this file.

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

import java.util.ArrayList;
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.UUID;

final class EmailMessageThreader {

    private int bogus_id_count = 0;
    
    private EmailMessageThreader(){}

    public static void threadMessages(List<EmailMessage> emailMessages) {
        EmailMessageThreader instance = new EmailMessageThreader();
        
        Map<String, EmailContainer> id_table = instance.createIDTable(emailMessages);
        Set<EmailContainer> rootSet = instance.getRootSet(id_table);

        instance.pruneEmptyContainers(rootSet);

        Set<EmailContainer> finalRootSet = instance.groupBySubject(rootSet);

        instance.assignThreadIDs(finalRootSet);
    }

    private Map<String, EmailContainer> createIDTable(List<EmailMessage> emailMessages) {
        HashMap<String, EmailContainer> id_table = new HashMap<>();

        for (EmailMessage message : emailMessages) {
            String messageID = message.getMessageID();

            // Check the id_table for an existing Container for message-id
            EmailContainer container = id_table.get(messageID);

            // An existing container for message-id was found
            if (container != null) {
                // If the existing Container has a message already assocated with it
                // emailMessage maybe a duplicate, so we don't lose the existance of 
                // the duplicate message assign it a bogus message-id
                if (container.hasMessage()) {
                    messageID = String.format("<Bogus-id: %d >", bogus_id_count++);
                    container = null;
                } else {
                    container.setMessage(message);
                }
            }

            if (container == null) {
                container = new EmailContainer(message);
                id_table.put(messageID, container);
            }

            processMessageReferences(message, container, id_table);
        }

        return id_table;
    }

    void processMessageReferences(EmailMessage message, EmailContainer container, Map<String, EmailContainer> id_table) {
        List<String> referenceList = message.getReferences();

        // Make sure the inReplyToID is in the list of references
        String inReplyToID = message.getInReplyToID();
        if (inReplyToID != null && !inReplyToID.isEmpty()) {
            if (referenceList == null) {
                referenceList = new ArrayList<>();
            }

            referenceList.add(inReplyToID);
        }

        // No references, nothing to do
        if (referenceList == null) {
            return;
        }

        EmailContainer parent_ref = null;
        EmailContainer ref;

        for (String refID : referenceList) {
            // Check id_table to see if there is already a container for this
            // reference id, if not create a new Container and add to table
            ref = id_table.get(refID);

            if (ref == null) {
                ref = new EmailContainer();
                id_table.put(refID, ref);
            }

            // Set the parent\child relationship between parent_ref and ref
            if (parent_ref != null
                    && !ref.hasParent()
                    && !parent_ref.equals(ref)
                    && !parent_ref.isChild(ref)) {
                ref.setParent(parent_ref);
                parent_ref.addChild(ref);
            }

            parent_ref = ref;
        }

        // If the parent_ref and container are already linked, don't change 
        // anything
        if (parent_ref != null
                && (parent_ref.equals(container)
                || container.isChild(parent_ref))) {
            parent_ref = null;
        }

        // If container already has a parent, the parent was assumed based on
        // the list of references from another message.  parent_ref will be
        // the real parent of container so throw away the old parent and set a
        // new one.
        if (container.hasParent()) {
            container.getParent().removeChild(container);
            container.setParent(null);
        }

        if (parent_ref != null) {
            container.setParent(container);
            parent_ref.addChild(container);
        }
    }

    Set<EmailContainer> getRootSet(Map<?, EmailContainer> id_table) {
        HashSet<EmailContainer> rootSet = new HashSet<>();

        id_table.values().stream().filter((container)
                -> (!container.hasParent())).forEachOrdered((container) -> {
            rootSet.add(container);
        });

        return rootSet;
    }

    void pruneEmptyContainers(Set<EmailContainer> containerSet) {
        Set<EmailContainer> containersToRemove = new HashSet<>();
        containerSet.forEach((container) -> {
            if (!container.hasMessage() && !container.hasChildren()) {
                containersToRemove.add(container);
            } else {
                pruneChildren(container);
            }
        });

        containerSet.removeAll(containersToRemove);
    }

    boolean pruneChildren(EmailContainer parent) {
        if (parent == null) {
            return false;
        }

        Set<EmailContainer> children = parent.getChildren();
        
        if (children == null) {
            return false;
        }

        EmailContainer grandParent = parent.getParent();
        Set<EmailContainer> remove = new HashSet<>();
        Set<EmailContainer> add = new HashSet<>();
        for (EmailContainer child : parent.getChildren()) {
            if (pruneChildren(child)) {
                remove.add(child);
                add.addAll(child.getChildren());
                child.setParent(null);
                child.clearChildren();

            }
        }

        parent.addChildren(add);
        parent.removeChildren(remove);

        if (!parent.hasMessage() && grandParent != null) {
            children.forEach((child) -> {
                child.setParent(grandParent);
            });
            return true;
        }

        return false;
    }

    Set<EmailContainer> groupBySubject(Set<EmailContainer> rootSet) {
        Map<String, EmailContainer> subject_table = createSubjectTable(rootSet);

        Set<EmailContainer> finalSet = new HashSet<>();

        for (EmailContainer rootSetContainer : rootSet) {
            String rootSubject = rootSetContainer.getSimplifiedSubject();

            EmailContainer tableContainer = subject_table.get(rootSubject);
            if (tableContainer == null || tableContainer.equals(rootSetContainer)) {
                finalSet.add(rootSetContainer);
                continue;
            }

            // If both containers are dummy/empty append the children of one to the other
            if (tableContainer.getMessage() == null && rootSetContainer.getMessage() == null) {
                tableContainer.addChildren(rootSetContainer.getChildren());
                rootSetContainer.clearChildren();
                continue;
            }

            // one container is empty, but the other is not, make the non-empty one be a
            // child of the empty
            if ((tableContainer.getMessage() == null && rootSetContainer.getMessage() != null)
                    || (tableContainer.getMessage() != null && rootSetContainer.getMessage() == null)) {

                if (tableContainer.getMessage() == null) {
                    tableContainer.addChild(rootSetContainer);

                } else {
                    rootSetContainer.addChild(tableContainer);
                    subject_table.remove(rootSubject, tableContainer);
                    subject_table.put(rootSubject, rootSetContainer);

                    finalSet.add(rootSetContainer);
                }

                continue;
            }

            // tableContainer is non-empty and it's message's subject does not begin
            // with 'RE:' but rootSetContainer's message does begin with 'RE:', then
            // make rootSetContainer a child of tableContainer
            if (tableContainer.getMessage() != null
                    && !tableContainer.isReplySubject()
                    && rootSetContainer.isReplySubject()) {
                tableContainer.addChild(rootSetContainer);
                continue;
            }

            // If table container is non-empy, and table container's subject does 
            // begin with 'RE:', but rootSetContainer does not start with 'RE:'
            // make tableContainer a child of rootSetContainer
            if (tableContainer.getMessage() != null
                    && tableContainer.isReplySubject()
                    && !rootSetContainer.isReplySubject()) {
                rootSetContainer.addChild(tableContainer);
                subject_table.put(rootSubject, rootSetContainer);
                finalSet.add(rootSetContainer);
                continue;
            }

            // rootSetContainer and tableContainer either both have 'RE' or
            // don't.  Create a new dummy container with both containers as
            // children.
            EmailContainer newParent = new EmailContainer();
            newParent.addChild(tableContainer);
            newParent.addChild(rootSetContainer);
            subject_table.remove(rootSubject, tableContainer);
            subject_table.put(rootSubject, newParent);

            finalSet.add(newParent);
        }
        return finalSet;
    }

    Map<String, EmailContainer> createSubjectTable(Set<EmailContainer> rootSet) {
        HashMap<String, EmailContainer> subject_table = new HashMap<>();

        for (EmailContainer rootSetContainer : rootSet) {
            String subject = "";
            boolean reSubject = false;

            if (rootSetContainer.hasMessage()) {
                subject = rootSetContainer.getMessage().getSimplifiedSubject();
                reSubject = rootSetContainer.getMessage().isReplySubject();
            } else if (rootSetContainer.hasChildren()) {
                Iterator<EmailContainer> childrenIterator = rootSetContainer.getChildren().iterator();
                while (childrenIterator.hasNext()) {
                    EmailMessage childMessage = childrenIterator.next().getMessage();
                    if (childMessage != null) {
                        subject = childMessage.getSimplifiedSubject();
                        if (!subject.isEmpty()) {
                            reSubject = childMessage.isReplySubject();
                            break;
                        }
                    }
                }
            }

            if (subject == null || subject.isEmpty()) {
                continue;  // Give up on this container
            }

            EmailContainer tableContainer = subject_table.get(subject);
            
//      A container will be added to the table, if a container for its "simplified" subject
//      does not currently exist in the table.  Or  if there is more than one container with the same
//      subject, but one is an "empty container" the empty one will be added 
//      the table or in the one in the table has "RE" in the subject it will be replaced
//      by the one that does not have "RE" in the subject (if it exists)
//     
            if (tableContainer == null || 
                    (tableContainer.getMessage() != null && rootSetContainer.getMessage() == null) ||
                    (!reSubject && (tableContainer.getMessage() != null && tableContainer.getMessage().isReplySubject()))) { 
                subject_table.put(subject, rootSetContainer);
            }

        }

        return subject_table;
    }
    
    private void assignThreadIDs(Set<EmailContainer> containerSet) {
        for(EmailContainer container: containerSet) {
            // Generate a threadID
            String threadID = UUID.randomUUID().toString();
            // Add the IDs to this thread
            addThreadID(container, threadID);
        }
    }
    
    private void addThreadID(EmailContainer container, String threadID) {
        if(container == null) {
            return;
        }
        
        EmailMessage message = container.getMessage();
        if(message != null) {
            message.setMessageThreadID(threadID);
        }
        
        if(container.hasChildren()) {
            for(EmailContainer child: container.getChildren()) {
                addThreadID(child, threadID);
            }
        }
    }

    final class EmailContainer {

        private EmailMessage message;
        private EmailContainer parent;
        private Set<EmailContainer> children;

        EmailContainer() {
            // This constructor is intentially empty to allow for the creation of 
            // an EmailContainer without a message
        }

        EmailContainer(EmailMessage message) {
            this.message = message;
        }

        EmailMessage getMessage() {
            return message;
        }

        void setMessage(EmailMessage message) {
            this.message = message;
        }

        boolean hasMessage() {
            return message != null;
        }

        String getSimplifiedSubject() {
            String subject = "";
            if (message != null) {
                subject = message.getSimplifiedSubject();
            } else if (children != null) {
                for (EmailContainer child : children) {
                    if (child.hasMessage()) {
                        subject = child.getSimplifiedSubject();
                    }

                    if (subject != null && !subject.isEmpty()) {
                        break;
                    }
                }
            }
            return subject;
        }

        boolean isReplySubject() {
            if (message != null) {
                return message.isReplySubject();
            } else if (children != null) {
                for (EmailContainer child : children) {
                    if (child.hasMessage()) {
                        boolean isReply = child.isReplySubject();

                        if (isReply) {
                            return isReply;
                        }
                    }
                }
            }

            return false;
        }

        EmailContainer getParent() {
            return parent;
        }

        void setParent(EmailContainer container) {
            parent = container;
        }

        boolean hasParent() {
            return parent != null;
        }

        boolean addChild(EmailContainer child) {
            if (children == null) {
                children = new HashSet<>();
            }

            return children.add(child);
        }

        boolean addChildren(Set<EmailContainer> children) {
            if (children == null || children.isEmpty()) {
                return false;
            }

            if (this.children == null) {
                this.children = new HashSet<>();
            }

            return this.children.addAll(children);
        }

        boolean removeChildren(Set<EmailContainer> children) {
            if (children != null) {
                return this.children.removeAll(children);
            }

            return false;
        }

        void clearChildren() {
            if( children != null ) {
                children.clear();
            }
        }

        boolean removeChild(EmailContainer child) {
            if(children != null) {
                return children.remove(child);
            } else {
                return false;
            }
        }

        boolean hasChildren() {
            return children != null && !children.isEmpty();
        }

        Set<EmailContainer> getChildren() {
            return children;
        }

        boolean isChild(EmailContainer container) {   
            return isChild(container, new HashSet<>());
        }   
        
        private boolean isChild(EmailContainer container, Set<EmailContainer> processedContainers) {
            processedContainers.add(this);
            if (children == null || children.isEmpty()) {
                return false;
            } else if (children.contains(container)) {
                return true;
            } else {
                for (EmailContainer child : children) {
                    // Prevent an infinite recursion by making sure we haven't already
                    // run isChild() on this child
                    if ((!processedContainers.contains(child)) && child.isChild(container, processedContainers)) {
                        return true;
                    }
                }
                return false;
            }
        }
    }
}