# -*- coding: utf-8 -*-
import psycopg2.extras
from datetime import datetime, timedelta
import traceback
import sys
import pandas as pd
import Utils
import logging
import logging.config
import yaml
import os
try:
import cPickle as pickle
except:
import pickle
import PredictionKorvesis as pdm
MYDIR = os.path.dirname(os.path.realpath(__file__))
LOGGING_CONF_FILE = os.path.join(MYDIR, "logging.yml")
def read_log_conf(yaml_file):
with open(yaml_file) as f:
logging.config.dictConfig(yaml.safe_load(f))
read_log_conf(LOGGING_CONF_FILE)
logger = logging.getLogger("mltd-offline")
def connect_time_db(
time_db_host,
time_db_port,
time_db_username,
time_db_password,
time_db_ssl,
time_db_dbname,
):
conn = psycopg2.connect(
dbname=time_db_dbname,
user=time_db_username,
password=time_db_password,
host=time_db_host,
)
cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
return cursor
def fetch_training_data(
time_db_cursor,
explicit_security_incidents_dates,
hours_before,
asset_id,
adt_table,
xlsiem_table,
od_table,
):
# fetch reported incidents from ADT table
time_db_cursor.execute(
"SELECT time FROM " + adt_table + " where asset_id = '" + asset_id + "'"
)
major_security_incidents_dates = time_db_cursor.fetchall()
# append to the fetched incident the explicit incidents
if len(explicit_security_incidents_dates) != 0:
for date_str in explicit_security_incidents_dates:
major_security_incidents_dates.append(
[datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%SZ"),"target_incident"]
)
major_security_incidents_dates = sorted(major_security_incidents_dates)
# removing incidents that are hours_before close to each other
remove_candidates = []
for i in range(len(major_security_incidents_dates) - 1, 0, -1):
cur_row = major_security_incidents_dates[i]
pre_row = major_security_incidents_dates[i - 1]
print(cur_row[0]-pre_row[0])
if cur_row[0] <= pre_row[0] + timedelta(hours=Utils.strtime_to_hours(hours_before)):
remove_candidates.append(cur_row)
if len(remove_candidates) > 0:
for date in remove_candidates:
major_security_incidents_dates.remove(date)
dataset_dates = []
dataset_values = []
for msi_date in major_security_incidents_dates:
time_db_cursor.execute(
"SELECT time, event_alarm_id FROM "
+ xlsiem_table
+ " where time > TIMESTAMP '"
+ msi_date[0].strftime("%Y-%m-%dT%H:%M:%SZ")
+ "' - interval '"
+ str(hours_before)
+ " hours' and asset_id = '"
+ asset_id
+ "'"
)
xlsiem_security_incidents = time_db_cursor.fetchall()
od_security_incidents = []
# time_db_cursor.execute(
# "SELECT time, event_alarm_id FROM "
# + od_table
# + " where time > TIMESTAMP '"
# + msi_date[0].strftime("%Y-%m-%dT%H:%M:%SZ")
# + "' - interval '"
# + str(hours_before)
# + " hours' and asset_id = '"
# + asset_id
# + "'"
# )
# od_security_incidents = time_db_cursor.fetchall()
fetched_security_incidents = sorted(xlsiem_security_incidents+od_security_incidents)
if len(fetched_security_incidents) > 0:
for sec_incident in fetched_security_incidents:
dataset_dates.append(sec_incident[0].strftime("%Y-%m-%dT%H:%M:%SZ"))
dataset_values.append(sec_incident[1])
dataset_dates.append(msi_date[0].strftime("%Y-%m-%dT%H:%M:%SZ"))
dataset_values.append("Exploits_All")
return (dataset_dates, dataset_values)
def save_to_bin_file(data, filename):
with open(filename, "ab") as f:
pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
def do_the_training(
trainID,
time_db_host,
time_db_port,
time_db_username,
time_db_password,
time_db_ssl,
time_db_dbname,
asset_id,
adt_table,
xlsiem_table,
od_table,
mp_thres_X,
mp_thres_Y,
mp_thres_Z,
pat_length,
s,
midpoint,
hours_before,
time_segments,
explicit_security_incidents_dates,
artificial_events_generation=False,
rre=False,
rfe=False,
kofe=False,
mil_over=False,
fs=False,
):
try:
logger.info("Training process started")
time_db_cursor = connect_time_db(
time_db_host,
time_db_port,
time_db_username,
time_db_password,
time_db_ssl,
time_db_dbname,
)
dataset_dates, dataset_values = fetch_training_data(
time_db_cursor,
explicit_security_incidents_dates,
hours_before,
asset_id,
adt_table,
xlsiem_table,
od_table,
)
target_event_id = "Exploits_All"
dataset = pd.DataFrame(
{"Timestamps": dataset_dates, "Event_id": dataset_values}
)
logger.info("serializing pattern length")
save_to_bin_file(pat_length, "train_" + str(trainID) + ".dat")
logger.info("serializing weak bins")
save_to_bin_file({}, "train_" + str(trainID) + ".dat")
logger.info("serializing matrixProfile")
save_to_bin_file({}, "train_" + str(trainID) + ".dat")
logger.info("serializing measurements")
save_to_bin_file(dataset_values, "train_" + str(trainID) + ".dat")
regr, feature_importance = pdm.train_rf_model(
dataset, time_segments, target_event_id, s, midpoint, rre=rre, rfe=rfe, kofe=kofe, mil_over=mil_over, fs=fs,
) # time_segments='H'
logger.info("serializing random forest trained model")
save_to_bin_file(regr, "train_" + str(trainID) + ".dat")
logger.info("serializing random forest feature importance")
save_to_bin_file(feature_importance, "train_" + str(trainID) + ".dat")
logger.info("serializing flag artificial_events_generation")
save_to_bin_file(artificial_events_generation, "train_" + str(trainID) + ".dat")
logger.info("Training process finished")
time_db_cursor.close()
return 0
except Exception as err:
print(err)
traceback.print_exc(file=sys.stdout)
return 5
# failure_dates = ["2014-05-07T00:01:02Z","2014-07-06T00:01:07Z","2014-10-03T00:01:00Z","2014-12-21T00:00:59Z","2015-02-14T00:00:59Z","2015-04-15T00:01:07Z","2015-07-01T00:00:45Z","2015-08-24T00:00:55Z","2016-04-24T00:00:50Z"]
# do_the_training(2,'localhost',8086,'Axoom3','artificial_events',0.5,1,0.5,6,0.7,2,4,failure_dates)