Dump database once at end

amrabed · Jun 9, 2015 · 5dc2c00 · 5dc2c00
1 parent 3d23e4b
commit 5dc2c00
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 19 deletions.
diff --git a/src/main/java/edu/vt/rhids/common/Database.java b/src/main/java/edu/vt/rhids/common/Database.java
@@ -1,12 +1,16 @@
 package edu.vt.rhids.common;
 
 import java.io.BufferedReader;
+import java.io.FileNotFoundException;
 import java.io.FileReader;
 import java.io.IOException;
+import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.HashMap;
 
 import edu.vt.rhids.input.BoSC;
+import edu.vt.rhids.util.Logger;
+import edu.vt.rhids.util.Logger.Verbosity;
 
 /**
  * Normal-behavior Database
@@ -102,6 +106,20 @@ public double calculateSimilarity(Database other)
 		return dot / (norm1 * norm2);
 	}
 
+	public void dump(String id) throws FileNotFoundException
+	{
+		final String file = "/var/log/rhids/db-" + id + ".dump";
+		try (PrintStream out = new PrintStream(file))
+		{
+			Logger.log("Dumping database to " + file, Verbosity.MEDIUM);
+			for (BoSC bosc : keySet())
+			{
+				out.println(bosc);
+			}
+			Logger.log("Done", Verbosity.MEDIUM);
+		}
+	}
+
 	@Override
 	public String toString()
 	{

diff --git a/src/main/java/edu/vt/rhids/main/Classifier.java b/src/main/java/edu/vt/rhids/main/Classifier.java
@@ -2,7 +2,6 @@
 
 import java.io.BufferedReader;
 import java.io.IOException;
-import java.io.PrintStream;
 import java.util.ArrayList;
 
 import edu.vt.rhids.common.Database;
@@ -41,30 +40,28 @@ public boolean trainUnconditionally() throws IOException
 		BoSC bosc;
 		String syscall;
 
-		try (PrintStream out = new PrintStream("/var/log/rhids/db-" + stats.getEpochSize() + ".dump"))
+		while (!RHIDS.isDoneTraining())
 		{
-			while (!RHIDS.isDoneTraining())
+			for (int i = 0; i < stats.getEpochSize(); i++)
 			{
-				for (int i = 0; i < stats.getEpochSize(); i++)
+				if ((syscall = SyscallParser.parse(reader)) != null)
 				{
-					if ((syscall = SyscallParser.parse(reader)) != null)
-					{
-						bosc = window.slide(syscall).getBoSC();
-						db.add(bosc);
-						out.println(bosc);
-						Logger.log(syscall + " => " + window + " => " + bosc, Verbosity.HIGH);
-					}
-					else
-					{
-						// Failed to train using input file
-						Logger.signal("Training failed", Verbosity.LOW);
-						return false;
-					}
+					bosc = window.slide(syscall).getBoSC();
+					db.add(bosc);
+					Logger.log(syscall + " => " + window + " => " + bosc, Verbosity.HIGH);
+				}
+				else
+				{
+					// Failed to train using input file
+					Logger.signal("Training failed", Verbosity.LOW);
+					db.dump(String.valueOf(stats.getEpochSize()));
+					return false;
 				}
-				Logger.log("Epoch " + stats.getTotalEpochs() + ": Database size is " + db.size(), Verbosity.MEDIUM);
-				stats.incrementTrainingEpochs();
 			}
+			Logger.log("Epoch " + stats.getTotalEpochs() + ": Database size is " + db.size(), Verbosity.MEDIUM);
+			stats.incrementTrainingEpochs();
 		}
+		db.dump(String.valueOf(stats.getEpochSize()));
 		return true;
 
 		// return train();