From c095514f9415ff904a86e9066637016b8e92befc Mon Sep 17 00:00:00 2001
From: BlackLight <blacklight@autistici.org>
Date: Tue, 26 Oct 2010 00:01:32 +0200
Subject: [PATCH] Weighted neural and bayesian networks correlation

---
 bayesian.c    | 23 +++++++++++++++++++++--
 correlation.c |  5 +++--
 neural.c      | 35 ++++++++++++++++++++++++++++++++++-
 spp_ai.h      |  5 +++++
 4 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/bayesian.c b/bayesian.c
index 55ea248..fec78ec 100644
--- a/bayesian.c
+++ b/bayesian.c
@@ -53,6 +53,20 @@ typedef struct  {
 PRIVATE AI_bayesian_correlation  *bayesian_cache    = NULL;
 PRIVATE double                   k_exp_value        = 0.0;
 
+/**
+ * \brief  Get the current weight of the bayesian correlation index using a hyperbolic tangent function with a parameter expressed in function of the current number of alerts in the history file
+ * \return The weight of the correlation index ( 0 <= weight < 1 )
+ */
+
+double
+AI_bayesian_correlation_weight ()
+{
+	double x = (double) AI_get_history_alert_number(),
+		  k = (double) config->alert_correlation_weight / HYPERBOLIC_TANGENT_SOLUTION;
+
+	return (( exp(x/k) - exp(-x/k) ) / ( exp(x/k) + exp(-x/k) ));
+}		/* -----  end of function AI_bayesian_correlation_weight  ----- */
+
 /**
  * \brief  Function used for computing the correlation probability A->B of two alerts (A,B) given their timestamps: f(ta, tb) = exp ( -(tb - ta)^2 / k )
  * \param  ta 	Timestamp of A
@@ -142,8 +156,13 @@ AI_alert_bayesian_correlation ( const AI_snort_alert *a, const AI_snort_alert *b
 			corr_count_a++;
 	}
 
-	corr /= (double) corr_count;
-	corr -= ( events_a->count - corr_count_a ) / events_a->count;
+	if ( !corr_count )
+	{
+		corr = 0.0;
+	} else {
+		corr /= (double) corr_count;
+		corr -= ( events_a->count - corr_count_a ) / events_a->count;
+	}
 
 	if ( found )
 	{
diff --git a/correlation.c b/correlation.c
index 4499261..7cfb252 100644
--- a/correlation.c
+++ b/correlation.c
@@ -1361,7 +1361,8 @@ AI_alert_correlation_thread ( void *arg )
 					/* Use the correlation indexes for which we have a value */
 					if ( bayesian_correlation != 0.0 && config->bayesianCorrelationInterval != 0 )
 					{
-						corr->correlation += bayesian_correlation;
+						corr->correlation += AI_bayesian_correlation_weight() * bayesian_correlation;
+						_dpd.logMsg ( "bayesian probability: %f\n", bayesian_correlation );
 						n_correlations++;
 					}
 
@@ -1373,7 +1374,7 @@ AI_alert_correlation_thread ( void *arg )
 
 					if ( neural_correlation != 0.0 && config->neuralNetworkTrainingInterval != 0 )
 					{
-						corr->correlation += neural_correlation;
+						corr->correlation += AI_neural_correlation_weight() * neural_correlation;
 						n_correlations++;
 					}
 
diff --git a/neural.c b/neural.c
index 690692d..1ef9c16 100644
--- a/neural.c
+++ b/neural.c
@@ -53,6 +53,39 @@ PRIVATE time_t latest_serialization_time  = ( time_t ) 0;
 PRIVATE som_network_t *net                = NULL;
 PRIVATE pthread_mutex_t neural_mutex;
 
+/**
+ * \brief  Get the current weight of the neural correlation index using a hyperbolic tangent function with a parameter expressed in function of the current number of alerts in the database
+ * \return The weight of the correlation index ( 0 <= weight < 1 )
+ */
+
+double
+AI_neural_correlation_weight ()
+{
+	DB_result res;
+	DB_row    row;
+	char      query[1024] = { 0 };
+	double    x = 0,
+			k = (double) config->alert_correlation_weight / HYPERBOLIC_TANGENT_SOLUTION;
+	
+	snprintf ( query, sizeof ( query ), "SELECT count(*) FROM %s", outdb_config[ALERTS_TABLE] );
+
+	if ( !DB_out_init() )
+	{
+		AI_fatal_err ( "Unable to connect to the database specified in module configuration", __FILE__, __LINE__ );
+	}
+
+	if ( !( res = (DB_result) DB_out_query ( query )))
+	{
+		AI_fatal_err ( "AIPreproc: Query error", __FILE__, __LINE__ );
+	}
+
+	row = (DB_row) DB_fetch_row ( res );
+	x = strtod ( row[0], NULL );
+	DB_free_result ( res );
+
+	return (( exp(x/k) - exp(-x/k) ) / ( exp(x/k) + exp(-x/k) ));
+}		/* -----  end of function AI_neural_correlation_weight  ----- */
+
 /**
  * \brief  Convert an alert row fetched from db to a vector suitable for being elaborated by the SOM neural network
  * \param  alert 	AI_som_alert_tuple object identifying the alert tuple
@@ -185,7 +218,7 @@ __AI_som_train ()
 	char      query[1024] = { 0 };
 
 	size_t    i = 0,
-			num_rows  = 0;
+			num_rows = 0;
 
 	DB_result res;
 	DB_row    row;
diff --git a/spp_ai.h b/spp_ai.h
index d54ee2e..9775192 100644
--- a/spp_ai.h
+++ b/spp_ai.h
@@ -105,6 +105,9 @@
 /** Cutoff y value in the exponential decay for considering two alerts not correlated */
 #define 	CUTOFF_Y_VALUE 					0.01
 
+/** Approximated solution of the equation tanh(x) = 0.95 */
+#define 	HYPERBOLIC_TANGENT_SOLUTION 	1.83178
+
 /****************************/
 /* Database support */
 #ifdef 	HAVE_LIBMYSQLCLIENT
@@ -497,6 +500,8 @@ const AI_alert_event*  AI_get_alert_events_by_key ( AI_alert_event_key );
 unsigned int           AI_get_history_alert_number ();
 double                 AI_alert_bayesian_correlation ( const AI_snort_alert*, const AI_snort_alert* );
 double                 AI_alert_neural_som_correlation ( const AI_snort_alert*, const AI_snort_alert* );
+double                 AI_neural_correlation_weight ();
+double                 AI_bayesian_correlation_weight ();
 
 void                   AI_outdb_mutex_initialize ();
 void*                  AI_store_alert_to_db_thread ( void* );