Merge pull request #5834 from gdicristofaro/6265-cvtVis

6265 cvt visualization message count changes

Core/src/org/sleuthkit/autopsy/communications/AbstractCVTAction.java

@@ -19,11 +19,13 @@
 package org.sleuthkit.autopsy.communications;
 
 import java.util.Collection;
+import java.util.stream.Collectors;
 import javax.swing.AbstractAction;
 import javax.swing.ImageIcon;
 import javax.swing.JMenuItem;
 import org.openide.util.Utilities;
 import org.openide.util.actions.Presenter;
+import org.sleuthkit.datamodel.AccountDeviceInstance;
 
 /**
  * Base class for actions that act on the selected AccountDeviceInstanceKeys.
@@ -37,8 +39,11 @@ abstract class AbstractCVTAction extends AbstractAction implements Presenter.Pop
      *
      * @return The selected accounts
      */
-    Collection<? extends AccountDeviceInstanceKey> getSelectedAccounts() {
+    Collection<? extends AccountDeviceInstance> getSelectedAccounts() {
-        return  Utilities.actionsGlobalContext().lookupAll(AccountDeviceInstanceKey.class);
+        return  Utilities.actionsGlobalContext().lookupAll(AccountDeviceInstanceKey.class)
+            .stream()
+            .map((adiKey -> adiKey.getAccountDeviceInstance()))
+            .collect(Collectors.toSet());
     }
 
     @Override

Core/src/org/sleuthkit/autopsy/communications/CVTEvents.java

@@ -24,6 +24,7 @@ import java.util.Collection;
 import org.sleuthkit.autopsy.communications.FiltersPanel.DateControlState;
 import org.sleuthkit.datamodel.CommunicationsFilter;
 import org.sleuthkit.autopsy.communications.StateManager.CommunicationsState;
+import org.sleuthkit.datamodel.AccountDeviceInstance;
 
 /**
  * Provide the singleton EventBus.
@@ -73,19 +74,19 @@ final class CVTEvents {
      */
     static final class PinAccountsEvent {
 
-        private final ImmutableSet<AccountDeviceInstanceKey> accountDeviceInstances;
+        private final ImmutableSet<AccountDeviceInstance> accounInstances;
         private final boolean replace;
 
         public boolean isReplace() {
             return replace;
         }
 
-        ImmutableSet<AccountDeviceInstanceKey> getAccountDeviceInstances() {
+        ImmutableSet<AccountDeviceInstance> getAccountDeviceInstances() {
-            return accountDeviceInstances;
+            return accounInstances;
         }
 
-        PinAccountsEvent(Collection<? extends AccountDeviceInstanceKey> accountDeviceInstances, boolean replace) {
+        PinAccountsEvent(Collection<? extends AccountDeviceInstance> accountDeviceInstances, boolean replace) {
-            this.accountDeviceInstances = ImmutableSet.copyOf(accountDeviceInstances);
+            this.accounInstances = ImmutableSet.copyOf(accountDeviceInstances);
             this.replace = replace;
         }
     }
@@ -95,14 +96,14 @@ final class CVTEvents {
      */
     static final class UnpinAccountsEvent {
 
-        private final ImmutableSet<AccountDeviceInstanceKey> accountDeviceInstances;
+        private final ImmutableSet<AccountDeviceInstance> accountInstances;
 
-        public ImmutableSet<AccountDeviceInstanceKey> getAccountDeviceInstances() {
+        public ImmutableSet<AccountDeviceInstance> getAccountDeviceInstances() {
-            return accountDeviceInstances;
+            return accountInstances;
         }
 
-         UnpinAccountsEvent(Collection<? extends AccountDeviceInstanceKey> accountDeviceInstances) {
+         UnpinAccountsEvent(Collection<? extends AccountDeviceInstance> accountDeviceInstances) {
-            this.accountDeviceInstances = ImmutableSet.copyOf(accountDeviceInstances);
+            this.accountInstances = ImmutableSet.copyOf(accountDeviceInstances);
         }
     }
     

Core/src/org/sleuthkit/autopsy/communications/CommunicationsGraph.java

@@ -32,6 +32,7 @@ import java.io.InputStreamReader;
 import java.io.StringWriter;
 import java.net.URL;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -147,7 +148,7 @@ final class CommunicationsGraph extends mxGraph {
             scopes.put("accountName", adiKey.getAccountDeviceInstance().getAccount().getTypeSpecificID());
             scopes.put("size", Math.round(Math.log(adiKey.getMessageCount()) + 5));
             scopes.put("iconFileName", CommunicationsGraph.class.getResource(Utils.getIconFilePath(adiKey.getAccountDeviceInstance().getAccount().getAccountType())));
-            scopes.put("pinned", pinnedAccountModel.isAccountPinned(adiKey));
+            scopes.put("pinned", pinnedAccountModel.isAccountPinned(adiKey.getAccountDeviceInstance()));
             scopes.put("MARKER_PIN_URL", MARKER_PIN_URL);
             scopes.put("locked", lockedVertexModel.isVertexLocked((mxCell) cell));
             scopes.put("LOCK_URL", LOCK_URL);
@@ -172,7 +173,7 @@ final class CommunicationsGraph extends mxGraph {
             scopes.put("accountName", adiKey.getAccountDeviceInstance().getAccount().getTypeSpecificID());
             scopes.put("relationships", 12);// Math.round(Math.log(adiKey.getMessageCount()) + 5));
             scopes.put("iconFileName", CommunicationsGraph.class.getResource(Utils.getIconFilePath(adiKey.getAccountDeviceInstance().getAccount().getAccountType())));
-            scopes.put("pinned", pinnedAccountModel.isAccountPinned(adiKey));
+            scopes.put("pinned", pinnedAccountModel.isAccountPinned(adiKey.getAccountDeviceInstance()));
             scopes.put("MARKER_PIN_URL", MARKER_PIN_URL);
             scopes.put("locked", lockedVertexModel.isVertexLocked((mxCell) cell));
             scopes.put("LOCK_URL", LOCK_URL);
@@ -199,16 +200,24 @@ final class CommunicationsGraph extends mxGraph {
         lockedVertexModel.clear();
     }
 
-    private mxCell getOrCreateVertex(AccountDeviceInstanceKey accountDeviceInstanceKey) {
+    private mxCell getOrCreateVertex(AccountDeviceInstance adi, CommunicationsManager commsManager, CommunicationsFilter currentFilter) {
-        final AccountDeviceInstance accountDeviceInstance = accountDeviceInstanceKey.getAccountDeviceInstance();
+        final AccountDeviceInstance accountDeviceInstance = adi;
         final String name = accountDeviceInstance.getAccount().getTypeSpecificID();
 
         final mxCell vertex = nodeMap.computeIfAbsent(name + accountDeviceInstance.getDeviceId(), vertexName -> {
-            double size = Math.sqrt(accountDeviceInstanceKey.getMessageCount()) + 10;
+            long adiRelationshipsCount = 1;
-
+            try {
+                adiRelationshipsCount = commsManager.getRelationshipSourcesCount(accountDeviceInstance, currentFilter);
+            } catch (TskCoreException tskCoreException) {
+                logger.log(Level.SEVERE, "There was an error fetching relationships for the node: " + accountDeviceInstance, tskCoreException);
+            }
+                        
+            double size = Math.sqrt(adiRelationshipsCount) + 10;
+            AccountDeviceInstanceKey adiKey = new AccountDeviceInstanceKey(adi, currentFilter, adiRelationshipsCount);
+            
             mxCell newVertex = (mxCell) insertVertex(
                     getDefaultParent(),
-                    name, accountDeviceInstanceKey,
+                    name, adiKey,
                     Math.random() * 400,
                     Math.random() * 400,
                     size,
@@ -219,9 +228,11 @@ final class CommunicationsGraph extends mxGraph {
     }
 
     @SuppressWarnings("unchecked")
-    private mxCell addOrUpdateEdge(long relSources, AccountDeviceInstanceKey account1, AccountDeviceInstanceKey account2) {
+    private mxCell addOrUpdateEdge(long relSources, 
-        mxCell vertex1 = getOrCreateVertex(account1);
+            AccountDeviceInstance account1, AccountDeviceInstance account2, 
-        mxCell vertex2 = getOrCreateVertex(account2);
+            CommunicationsManager commsManager, CommunicationsFilter currentFilter) {
+        mxCell vertex1 = getOrCreateVertex(account1, commsManager, currentFilter);
+        mxCell vertex2 = getOrCreateVertex(account2, commsManager, currentFilter);
         Object[] edgesBetween = getEdgesBetween(vertex1, vertex2);
         mxCell edge;
         if (edgesBetween.length == 0) {
@@ -260,28 +271,24 @@ final class CommunicationsGraph extends mxGraph {
                 /**
                  * set to keep track of accounts related to pinned accounts
                  */
-                final Map<AccountDeviceInstance, AccountDeviceInstanceKey> relatedAccounts = new HashMap<>();
+                final Set<AccountDeviceInstance> relatedAccounts = new HashSet<>();
-                for (final AccountDeviceInstanceKey adiKey : pinnedAccountModel.getPinnedAccounts()) {
+                for (final AccountDeviceInstance adi : pinnedAccountModel.getPinnedAccounts()) {
                     if (isCancelled()) {
                         break;
                     }
                     //get accounts related to pinned account
                     final List<AccountDeviceInstance> relatedAccountDeviceInstances
-                            = commsManager.getRelatedAccountDeviceInstances(adiKey.getAccountDeviceInstance(), currentFilter);
+                            = commsManager.getRelatedAccountDeviceInstances(adi, currentFilter);
-                    relatedAccounts.put(adiKey.getAccountDeviceInstance(), adiKey);
+                    relatedAccounts.add(adi);
-                    getOrCreateVertex(adiKey);
+                    getOrCreateVertex(adi, commsManager, currentFilter);
+
+                    for (final AccountDeviceInstance relatedADI : relatedAccountDeviceInstances)
+                        relatedAccounts.add(relatedADI);
 
-                    for (final AccountDeviceInstance relatedADI : relatedAccountDeviceInstances) {
-                        final long adiRelationshipsCount = commsManager.getRelationshipSourcesCount(relatedADI, currentFilter);
-                        final AccountDeviceInstanceKey relatedADIKey = new AccountDeviceInstanceKey(relatedADI, currentFilter, adiRelationshipsCount);
-                        relatedAccounts.put(relatedADI, relatedADIKey); //store related accounts
-                    }
                     progressIndicator.progress(++progressCounter);
                 }
 
-                Set<AccountDeviceInstance> accounts = relatedAccounts.keySet();
+                Map<AccountPair, Long> relationshipCounts = commsManager.getRelationshipCountsPairwise(relatedAccounts, currentFilter);
-
-                Map<AccountPair, Long> relationshipCounts = commsManager.getRelationshipCountsPairwise(accounts, currentFilter);
 
                 int total = relationshipCounts.size();
                 int progress = 0;
@@ -290,12 +297,12 @@ final class CommunicationsGraph extends mxGraph {
                 for (Map.Entry<AccountPair, Long> entry : relationshipCounts.entrySet()) {
                     Long count = entry.getValue();
                     AccountPair relationshipKey = entry.getKey();
-                    AccountDeviceInstanceKey account1 = relatedAccounts.get(relationshipKey.getFirst());
+                    AccountDeviceInstance account1 = relationshipKey.getFirst();
-                    AccountDeviceInstanceKey account2 = relatedAccounts.get(relationshipKey.getSecond());
+                    AccountDeviceInstance account2 = relationshipKey.getSecond();
 
                     if (pinnedAccountModel.isAccountPinned(account1)
                             || pinnedAccountModel.isAccountPinned(account2)) {
-                        mxCell addEdge = addOrUpdateEdge(count, account1, account2);
+                        mxCell addEdge = addOrUpdateEdge(count, account1, account2, commsManager, currentFilter);
                         progressText = addEdge.getId();
                     }
                     progressIndicator.progress(progressText, progress++);

Core/src/org/sleuthkit/autopsy/communications/PinnedAccountModel.java

@@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.eventbus.EventBus;
 import java.util.HashSet;
 import java.util.Set;
+import org.sleuthkit.datamodel.AccountDeviceInstance;
 
 /**
  * Model of what accounts are pinned to a visualization.
@@ -34,7 +35,7 @@ class PinnedAccountModel {
      * to pinned accounts and pass the filters are show. Pinning accounts is the
      * primary way to populate the graph.
      */
-    private final Set<AccountDeviceInstanceKey> pinnedAccountDevices = new HashSet<>();
+    private final Set<AccountDeviceInstance> pinnedAccountDevices = new HashSet<>();
 
     private final EventBus eventBus = new EventBus();
 
@@ -46,7 +47,7 @@ class PinnedAccountModel {
         eventBus.unregister(handler);
     }
 
-    boolean isAccountPinned(AccountDeviceInstanceKey account) {
+    boolean isAccountPinned(AccountDeviceInstance account) {
         return pinnedAccountDevices.contains(account);
     }
 
@@ -57,7 +58,7 @@ class PinnedAccountModel {
      *
      * @param accountDeviceInstances The accounts to unpin.
      */
-    void unpinAccount(Set<AccountDeviceInstanceKey> accountDeviceInstances) {
+    void unpinAccount(Set<AccountDeviceInstance> accountDeviceInstances) {
         pinnedAccountDevices.removeAll(accountDeviceInstances);
     }
 
@@ -68,7 +69,7 @@ class PinnedAccountModel {
      *
      * @param accountDeviceInstances The accounts to pin.
      */
-    void pinAccount(Set<AccountDeviceInstanceKey> accountDeviceInstances) {
+    void pinAccount(Set<AccountDeviceInstance> accountDeviceInstances) {
         pinnedAccountDevices.addAll(accountDeviceInstances);
     }
 
@@ -86,7 +87,7 @@ class PinnedAccountModel {
         pinnedAccountDevices.clear();
     }
 
-    ImmutableSet<AccountDeviceInstanceKey> getPinnedAccounts() {
+    ImmutableSet<AccountDeviceInstance> getPinnedAccounts() {
         return ImmutableSet.copyOf(pinnedAccountDevices);
     }
 

Core/src/org/sleuthkit/autopsy/communications/StateManager.java

@@ -23,6 +23,7 @@ import java.util.HashSet;
 import java.util.Set;
 import org.sleuthkit.autopsy.communications.FiltersPanel.DateControlState;
 import org.sleuthkit.autopsy.coreutils.History;
+import org.sleuthkit.datamodel.AccountDeviceInstance;
 import org.sleuthkit.datamodel.CommunicationsFilter;
 
 /**
@@ -51,11 +52,11 @@ final class StateManager {
     @Subscribe
     void pinAccount(CVTEvents.PinAccountsEvent pinEvent) {
         if(pinEvent.isReplace()){
-            HashSet<AccountDeviceInstanceKey> pinnedList = new HashSet<>();
+            HashSet<AccountDeviceInstance> pinnedList = new HashSet<>();
             pinnedList.addAll(pinEvent.getAccountDeviceInstances());
             historyManager.advance(new CommunicationsState(comFilter, pinnedList, -1, currentStartState, currentEndState));
         } else {
-            HashSet<AccountDeviceInstanceKey> pinnedList = new HashSet<>();
+            HashSet<AccountDeviceInstance> pinnedList = new HashSet<>();
             pinnedList.addAll(pinEvent.getAccountDeviceInstances());
             pinnedList.addAll(pinModel.getPinnedAccounts());
             
@@ -74,7 +75,7 @@ final class StateManager {
     @Subscribe
     void unpinAccounts(CVTEvents.UnpinAccountsEvent pinEvent) {
 
-        HashSet<AccountDeviceInstanceKey> pinnedList = new HashSet<>();
+        HashSet<AccountDeviceInstance> pinnedList = new HashSet<>();
         pinnedList.addAll(pinModel.getPinnedAccounts());
         pinnedList.removeAll(pinEvent.getAccountDeviceInstances());
         
@@ -135,7 +136,7 @@ final class StateManager {
      */
     final class CommunicationsState{
         private final CommunicationsFilter communcationFilter;
-        private final Set<AccountDeviceInstanceKey> pinnedList;
+        private final Set<AccountDeviceInstance> pinnedList;
         private final double zoomValue;
         private final DateControlState startDateState;
         private final DateControlState endDateState;
@@ -149,7 +150,7 @@ final class StateManager {
          * @param zoomValue Double value of the current graph scale
          */
         protected CommunicationsState(CommunicationsFilter communcationFilter, 
-                Set<AccountDeviceInstanceKey> pinnedList, double zoomValue, 
+                Set<AccountDeviceInstance> pinnedList, double zoomValue, 
                 DateControlState startDateState, DateControlState endDateState){
             this.pinnedList = pinnedList;
             this.communcationFilter = communcationFilter;
@@ -172,7 +173,7 @@ final class StateManager {
          * 
          * @return Set of AccountDeviceInstanceKey
          */
-        public Set<AccountDeviceInstanceKey> getPinnedList(){
+        public Set<AccountDeviceInstance> getPinnedList(){
             return pinnedList;
         }
 

Core/src/org/sleuthkit/autopsy/communications/VisualizationPanel.java

@@ -1100,7 +1100,7 @@ final public class VisualizationPanel extends JPanel {
                     } else {
                         jPopupMenu.add(new JMenuItem(new LockAction(selectedVertices)));
                     }
-                    if (pinnedAccountModel.isAccountPinned(adiKey)) {
+                    if (pinnedAccountModel.isAccountPinned(adiKey.getAccountDeviceInstance())) {
                         jPopupMenu.add(UnpinAccountsAction.getInstance().getPopupPresenter());
                     } else {
                         jPopupMenu.add(PinAccountsAction.getInstance().getPopupPresenter());