Refactored the database import workers

Core/src/org/sleuthkit/autopsy/modules/hashdatabase/ImportCentralRepoDbProgressDialog.java

@@ -111,7 +111,7 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
         return null;
     }
     
-    @NbBundle.Messages({"ImportCentralRepoDbProgressDialog.linesProcessed= lines processed"})
+    @NbBundle.Messages({"ImportCentralRepoDbProgressDialog.linesProcessed= hashes processed"})
     @Override
     public void propertyChange(PropertyChangeEvent evt) {
         
@@ -139,23 +139,23 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
         return worker.getLinesProcessed() + Bundle.ImportCentralRepoDbProgressDialog_linesProcessed();
     }
     
-    abstract class CentralRepoImportWorker2 extends SwingWorker<Void, Void>{
+    abstract class CentralRepoImportWorker extends SwingWorker<Void, Void>{
-        private final int HASH_IMPORT_THRESHOLD = 10000;
+        final int HASH_IMPORT_THRESHOLD = 10000;
-        private final String hashSetName;
+        final String hashSetName;
-        private final String version;
+        final String version;
-        private final int orgId;
+        final int orgId;
-        private final boolean searchDuringIngest;
+        final boolean searchDuringIngest;
-        private final boolean sendIngestMessages;
+        final boolean sendIngestMessages;
-        private final HashDbManager.HashDb.KnownFilesType knownFilesType;
+        final HashDbManager.HashDb.KnownFilesType knownFilesType;
-        private final boolean readOnly;
+        final boolean readOnly;
-        private final File importFile;
+        final File importFile;
-        private long totalLines;
+        long totalHashes = 1;
-        private int crIndex = -1;
+        int referenceSetID = -1;
-        private HashDbManager.CentralRepoHashDb newHashDb = null;
+        HashDbManager.CentralRepoHashDb newHashDb = null;
-        private final AtomicLong numLines = new AtomicLong();
+        final AtomicLong numLines = new AtomicLong();
-        private String errorString = "";
+        String errorString = "";
         
-        CentralRepoImportWorker2(String hashSetName, String version, int orgId,
+        CentralRepoImportWorker(String hashSetName, String version, int orgId,
             boolean searchDuringIngest, boolean sendIngestMessages, HashDbManager.HashDb.KnownFilesType knownFilesType,
             boolean readOnly, File importFile){
             
@@ -186,88 +186,88 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
             return errorString;
         }
         
+        /**
+         * Should be called in the constructor to set the max number of hashes.
+         * The value can be updated later after parsing the import file.
+         */
+        abstract void setEstimatedTotalHashes();
+        
+        void deleteIncompleteSet(){
+            if(referenceSetID >= 0){
+                
+                // This can be slow on large reference sets
+                Executors.newSingleThreadExecutor().execute(new Runnable() {
+                    @Override 
+                    public void run() {
+                        try{
+                            EamDb.getInstance().deleteReferenceSet(referenceSetID);
+                        } catch (EamDbException ex2){
+                            Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error deleting incomplete hash set from central repository", ex2);
+                        }
+                    }
+                });
+            }
         }
         
-    private interface CentralRepoImportWorker{
+        @NbBundle.Messages({"ImportCentralRepoDbProgressDialog.importHashsetError=Error importing hash set",
+                            "ImportCentralRepoDbProgressDialog.addDbError.message=Error adding new hash set"})
+        @Override
+        protected void done() {
             
-        void execute();
+            if(isCancelled()){
-        boolean cancel(boolean mayInterruptIfRunning);
+                // If the user hit cancel, delete this incomplete hash set from the central repo
-        void addPropertyChangeListener(PropertyChangeListener dialog);
+                deleteIncompleteSet();
-        int getProgressPercentage();
+                return;
-        long getLinesProcessed();
-        HashDbManager.HashDatabase getDatabase();
-        String getError();
             }
             
-    class ImportEncaseWorker extends SwingWorker<Void,Void> implements CentralRepoImportWorker{
+            try {
-        private final int HASH_IMPORT_THRESHOLD = 10000;
+                get();
-        private final String hashSetName;
+                try{
-        private final String version;
+                    newHashDb = HashDbManager.getInstance().addExistingCentralRepoHashSet(hashSetName, version, 
-        private final int orgId;
+                            referenceSetID, 
-        private final boolean searchDuringIngest;
+                            searchDuringIngest, sendIngestMessages, knownFilesType, readOnly);
-        private final boolean sendIngestMessages;
+                } catch (TskCoreException ex){
-        private final HashDbManager.HashDb.KnownFilesType knownFilesType;
+                    JOptionPane.showMessageDialog(null, Bundle.ImportCentralRepoDbProgressDialog_addDbError_message());
-        private final boolean readOnly;
+                    Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error adding imported hash set", ex);
-        private final File importFile;
+                }
-        private final long totalLines;
+            } catch (Exception ex) {
-        private int crIndex = -1;
+                // Delete this incomplete hash set from the central repo
-        private HashDbManager.CentralRepoHashDb newHashDb = null;
+                deleteIncompleteSet();
-        private final AtomicLong numLines = new AtomicLong();
+                errorString = Bundle.ImportCentralRepoDbProgressDialog_importHashsetError();
-        private String errorString = "";
+            }
+        }       
+                
+    }
+    
+    class ImportEncaseWorker extends CentralRepoImportWorker{
         
         ImportEncaseWorker(String hashSetName, String version, int orgId,
                 boolean searchDuringIngest, boolean sendIngestMessages, HashDbManager.HashDb.KnownFilesType knownFilesType,
                 boolean readOnly, File importFile){
+            super(hashSetName, version, orgId, searchDuringIngest, sendIngestMessages, knownFilesType, readOnly, importFile);
             
-            this.hashSetName = hashSetName;
+            setEstimatedTotalHashes();
-            this.version = version;
-            this.orgId = orgId;
-            this.searchDuringIngest = searchDuringIngest;
-            this.sendIngestMessages = sendIngestMessages;
-            this.knownFilesType = knownFilesType;
-            this.readOnly = readOnly;
-            this.importFile = importFile;
-            this.numLines.set(0);
-            
-            this.totalLines = getEstimatedTotalHashes();
         }
         
+        
         /**
          * Encase files have a 0x480 byte header, then each hash takes 18 bytes
          * @return Approximate number of hashes in the file
          */
-        final long getEstimatedTotalHashes(){
+        @Override
+        final void setEstimatedTotalHashes(){
             long fileSize = importFile.length();
             if(fileSize < 0x492){
-                return 1; // There's room for at most one hash
+                totalHashes =  1; // There's room for at most one hash
             }
-            return ((fileSize - 0x492) / 18);
+            totalHashes = (fileSize - 0x492) / 18;
-        }
-        
-        @Override 
-        public HashDbManager.HashDatabase getDatabase(){
-            return newHashDb;
-        }
-        
-        @Override
-        public long getLinesProcessed(){
-            return numLines.get();
-        }
-        
-        @Override
-        public int getProgressPercentage(){
-            return this.getProgress();
-        }
-        
-        @Override
-        public String getError(){
-            return errorString;
         }
         
         @Override
         protected Void doInBackground() throws Exception {
 
-            EncaseHashSetParser encaseParser = new EncaseHashSetParser(this.importFile.getAbsolutePath());
+            EncaseHashSetParser encaseParser = new EncaseHashSetParser(importFile.getAbsolutePath());
+            totalHashes = encaseParser.getExpectedHashes();
             
             TskData.FileKnown knownStatus;
             if (knownFilesType.equals(HashDbManager.HashDb.KnownFilesType.KNOWN)) {
@@ -277,7 +277,7 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
             }
             
             // Create an empty hashset in the central repository
-            crIndex = EamDb.getInstance().newReferenceSet(orgId, hashSetName, version, knownStatus, readOnly);
+            referenceSetID = EamDb.getInstance().newReferenceSet(orgId, hashSetName, version, knownStatus, readOnly);
 
             EamDb dbManager = EamDb.getInstance();
             CorrelationAttribute.Type contentType = dbManager.getCorrelationTypeById(CorrelationAttribute.FILES_TYPE_ID); // get "FILES" type
@@ -293,7 +293,7 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
 
                 if(newHash != null){
                     EamGlobalFileInstance eamGlobalFileInstance = new EamGlobalFileInstance(
-                            crIndex, 
+                            referenceSetID, 
                             newHash, 
                             knownStatus, 
                             "");
@@ -305,7 +305,7 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
                         dbManager.bulkInsertReferenceTypeEntries(globalInstances, contentType);
                         globalInstances.clear();
 
-                        int progress = (int)(numLines.get() * 100 / totalLines);
+                        int progress = (int)(numLines.get() * 100 / totalHashes);
                         if(progress < 100){
                             this.setProgress(progress);
                         } else {
@@ -319,90 +319,17 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
             this.setProgress(100);
             return null;
         }
-        
-        private void deleteIncompleteSet(int crIndex){
-            if(crIndex >= 0){
-                
-                // This can be slow on large reference sets
-                Executors.newSingleThreadExecutor().execute(new Runnable() {
-                    @Override 
-                    public void run() {
-                        try{
-                            EamDb.getInstance().deleteReferenceSet(crIndex);
-                        } catch (EamDbException ex2){
-                            Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error deleting incomplete hash set from central repository", ex2);
-                        }
-                    }
-                });
-            }
     }
     
-        @NbBundle.Messages({"ImportCentralRepoDbProgressDialog.importError=Error importing hash set"})
-        @Override
-        protected void done() {
     
-            if(isCancelled()){
+    class ImportIDXWorker extends CentralRepoImportWorker{
-                // If the user hit cancel, delete this incomplete hash set from the central repo
-                deleteIncompleteSet(crIndex);
-                return;
-            }
-            
-            try {
-                get();
-                try{
-                    newHashDb = HashDbManager.getInstance().addExistingCentralRepoHashSet(hashSetName, version, 
-                            crIndex, 
-                            searchDuringIngest, sendIngestMessages, knownFilesType, readOnly);
-                } catch (TskCoreException ex){
-                    JOptionPane.showMessageDialog(null, Bundle.ImportCentralRepoDbProgressDialog_addDbError_message());
-                    Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error adding imported hash set", ex);
-                }
-            } catch (Exception ex) {
-                // Delete this incomplete hash set from the central repo
-                if(crIndex >= 0){
-                    try{
-                        EamDb.getInstance().deleteReferenceSet(crIndex);
-                    } catch (EamDbException ex2){
-                        Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error deleting incomplete hash set from central repository", ex);
-                    }
-                }
-                errorString = Bundle.ImportCentralRepoDbProgressDialog_importError();
-            }
-        }       
-    }
-    
-    class ImportIDXWorker extends SwingWorker<Void,Void> implements CentralRepoImportWorker{
-        
-        private final int HASH_IMPORT_THRESHOLD = 10000;
-        private final String hashSetName;
-        private final String version;
-        private final int orgId;
-        private final boolean searchDuringIngest;
-        private final boolean sendIngestMessages;
-        private final HashDbManager.HashDb.KnownFilesType knownFilesType;
-        private final boolean readOnly;
-        private final File importFile;
-        private final long totalLines;
-        private int referenceSetID = -1;
-        private HashDbManager.CentralRepoHashDb newHashDb = null;
-        private final AtomicLong numLines = new AtomicLong();
-        private String errorString = "";
         
         ImportIDXWorker(String hashSetName, String version, int orgId,
                 boolean searchDuringIngest, boolean sendIngestMessages, HashDbManager.HashDb.KnownFilesType knownFilesType,
                 boolean readOnly, File importFile){
+            super(hashSetName, version, orgId, searchDuringIngest, sendIngestMessages, knownFilesType, readOnly, importFile);
             
-            this.hashSetName = hashSetName;
+            setEstimatedTotalHashes();
-            this.version = version;
-            this.orgId = orgId;
-            this.searchDuringIngest = searchDuringIngest;
-            this.sendIngestMessages = sendIngestMessages;
-            this.knownFilesType = knownFilesType;
-            this.readOnly = readOnly;
-            this.importFile = importFile;
-            this.numLines.set(0);
-            
-            this.totalLines = getEstimatedTotalHashes();
         }
         
         /**
@@ -411,29 +338,10 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
          * progress bar.
          * @return Approximate number of hashes in the file
          */
-        final long getEstimatedTotalHashes(){
+        @Override
+        final void setEstimatedTotalHashes(){
             long fileSize = importFile.length();
-            return (fileSize / 0x33 + 1); // IDX file lines are generally 0x33 bytes long, and we don't want this to be zero
+            totalHashes = fileSize / 0x33 + 1; // IDX file lines are generally 0x33 bytes long, and we don't want this to be zero
-        }
-        
-        @Override 
-        public HashDbManager.HashDatabase getDatabase(){
-            return newHashDb;
-        }
-        
-        @Override
-        public long getLinesProcessed(){
-            return numLines.get();
-        }
-        
-        @Override
-        public int getProgressPercentage(){
-            return this.getProgress();
-        }
-        
-        @Override
-        public String getError(){
-            return errorString;
         }
         
         @Override
@@ -480,7 +388,7 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
                     dbManager.bulkInsertReferenceTypeEntries(globalInstances, contentType);
                     globalInstances.clear();
 
-                    int progress = (int)(numLines.get() * 100 / totalLines);
+                    int progress = (int)(numLines.get() * 100 / totalHashes);
                     if(progress < 100){
                         this.setProgress(progress);
                     } else {
@@ -494,54 +402,6 @@ class ImportCentralRepoDbProgressDialog extends javax.swing.JDialog implements P
 
             return null;
         }
-        
-        private void deleteIncompleteSet(int idToDelete){
-            if(idToDelete >= 0){
-                
-                // This can be slow on large reference sets
-                Executors.newSingleThreadExecutor().execute(new Runnable() {
-                    @Override 
-                    public void run() {
-                        try{
-                            EamDb.getInstance().deleteReferenceSet(idToDelete);
-                        } catch (EamDbException ex2){
-                            Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error deleting incomplete hash set from central repository", ex2);
-                        }
-                    }
-                });
-            }
-        }
-        
-        @NbBundle.Messages({"ImportCentralRepoDbProgressDialog.addDbError.message=Error adding new hash set"})
-        @Override
-        protected void done() {
-            if(isCancelled()){
-                // If the user hit cancel, delete this incomplete hash set from the central repo
-                deleteIncompleteSet(referenceSetID);
-                return;
-            }
-            
-            try {
-                get();
-                try{
-                    newHashDb = HashDbManager.getInstance().addExistingCentralRepoHashSet(hashSetName, version, 
-                            referenceSetID, 
-                            searchDuringIngest, sendIngestMessages, knownFilesType, readOnly);
-                } catch (TskCoreException ex){
-                    JOptionPane.showMessageDialog(null, Bundle.ImportCentralRepoDbProgressDialog_addDbError_message());
-                    Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error adding imported hash set", ex);
-                }
-            } catch (Exception ex) {
-                // Delete this incomplete hash set from the central repo
-                if(referenceSetID >= 0){
-                    try{
-                        EamDb.getInstance().deleteReferenceSet(referenceSetID);
-                    } catch (EamDbException ex2){
-                        Logger.getLogger(ImportCentralRepoDbProgressDialog.class.getName()).log(Level.SEVERE, "Error deleting incomplete hash set from central repository", ex);
-                    }
-                }
-            }
-        }       
     }
     
     /**