From 820135d7723c229a1028ace0e458d3852f018310 Mon Sep 17 00:00:00 2001
From: Brian Carrier <carrier@sleuthkit.org>
Date: Thu, 11 Jun 2015 00:40:07 -0400
Subject: [PATCH] Created Keys class for DataSource to update more cleanly when
 new data sources are added

---
 .../datamodel/AbstractContentChildren.java    |  6 +-
 .../autopsy/datamodel/DataSourcesNode.java    | 58 ++++++++++++++++++-
 .../DirectoryTreeTopComponent.java            |  3 +-
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/Core/src/org/sleuthkit/autopsy/datamodel/AbstractContentChildren.java b/Core/src/org/sleuthkit/autopsy/datamodel/AbstractContentChildren.java
index 77ca0752fb..0719a18c3c 100644
--- a/Core/src/org/sleuthkit/autopsy/datamodel/AbstractContentChildren.java
+++ b/Core/src/org/sleuthkit/autopsy/datamodel/AbstractContentChildren.java
@@ -176,11 +176,7 @@ abstract class AbstractContentChildren<T> extends Keys<T> {
 
         @Override
         public AbstractNode visit(DataSources i) {
-            try {
-                return new DataSourcesNode(Case.getCurrentCase().getDataSources());
-            } catch (TskCoreException ex) {
-                return defaultVisit(i);
-            }
+            return new DataSourcesNode();
         }
 
         @Override
diff --git a/Core/src/org/sleuthkit/autopsy/datamodel/DataSourcesNode.java b/Core/src/org/sleuthkit/autopsy/datamodel/DataSourcesNode.java
index 4d0ed1feb4..1d45eec0dc 100644
--- a/Core/src/org/sleuthkit/autopsy/datamodel/DataSourcesNode.java
+++ b/Core/src/org/sleuthkit/autopsy/datamodel/DataSourcesNode.java
@@ -18,26 +18,78 @@
  */
 package org.sleuthkit.autopsy.datamodel;
 
+import java.beans.PropertyChangeEvent;
+import java.beans.PropertyChangeListener;
+import java.util.Collections;
 import java.util.List;
-
 import org.openide.nodes.Sheet;
 import org.openide.util.NbBundle;
 import org.openide.util.lookup.Lookups;
+import org.sleuthkit.autopsy.casemodule.Case;
+import org.sleuthkit.autopsy.coreutils.Logger;
 import org.sleuthkit.datamodel.Content;
+import org.sleuthkit.datamodel.TskCoreException;
 
 /**
  * Nodes for the images
  */
 public class DataSourcesNode extends DisplayableItemNode {
-
     public static final String NAME = NbBundle.getMessage(DataSourcesNode.class, "DataSourcesNode.name");
 
+    // NOTE: The images passed in via argument will be ignored.
+    @Deprecated
     public DataSourcesNode(List<Content> images) {
-        super(new RootContentChildren(images), Lookups.singleton(NAME));
+        super(new DataSourcesNodeChildren(), Lookups.singleton(NAME));
+        init();
+    }
+    
+    public DataSourcesNode() {
+        super(new DataSourcesNodeChildren(), Lookups.singleton(NAME));
+        init();
+    }
+    
+    private void init() {    
         setName(NAME);
         setDisplayName(NAME);
         this.setIconBaseWithExtension("org/sleuthkit/autopsy/images/image.png"); //NON-NLS
     }
+    
+    /* Custom Keys implementation that listens for new data sources being added. */
+    private static class DataSourcesNodeChildren extends AbstractContentChildren<Content> {
+        private static final Logger logger = Logger.getLogger(DataSourcesNodeChildren.class.getName());
+        
+        public DataSourcesNodeChildren() {
+            super();
+        }
+
+        private final PropertyChangeListener pcl = new PropertyChangeListener() {
+            @Override
+            public void propertyChange(PropertyChangeEvent evt) {
+                String eventType = evt.getPropertyName();
+                if (eventType.equals(Case.Events.DATA_SOURCE_ADDED.toString())) {
+                    addNotify();
+                }
+            }
+        };
+        
+
+        @Override
+        protected void addNotify() {
+            Case.addPropertyChangeListener(pcl);
+            try {
+                setKeys(Case.getCurrentCase().getDataSources());
+            } catch (TskCoreException | IllegalStateException ex) {
+                logger.severe("Error getting data sources: " + ex.getMessage()); // NON-NLS
+                setKeys(Collections.EMPTY_SET);
+            }
+        }
+
+        @Override
+        protected void removeNotify() {
+            Case.removePropertyChangeListener(pcl);
+            setKeys(Collections.EMPTY_SET);
+        }
+    }
 
     @Override
     public boolean isLeafTypeNode() {
diff --git a/Core/src/org/sleuthkit/autopsy/directorytree/DirectoryTreeTopComponent.java b/Core/src/org/sleuthkit/autopsy/directorytree/DirectoryTreeTopComponent.java
index 8161914bb9..c22515dfe2 100644
--- a/Core/src/org/sleuthkit/autopsy/directorytree/DirectoryTreeTopComponent.java
+++ b/Core/src/org/sleuthkit/autopsy/directorytree/DirectoryTreeTopComponent.java
@@ -573,7 +573,8 @@ public final class DirectoryTreeTopComponent extends TopComponent implements Dat
             }
         } // if the image is added to the case
         else if (changed.equals(Case.Events.DATA_SOURCE_ADDED.toString())) {
-            componentOpened();
+            // we don't need to do anything in here. 
+            // DataSourcesNode is listening for these events and updates itself
         }
         // change in node selection
         else if (changed.equals(ExplorerManager.PROP_SELECTED_NODES)) {