def construct_port_summary_list(self, rqst, uom, units):
g_config = GlobalConfig()
plugins_dictionary = g_config.get_plugin_dictionary()
json_list = []
for key, val in plugins_dictionary.items():
json_list.append(self.construct_port_summary(rqst, val['port'], val['table'], uom, units))
return json_list
class Database:
def __init__(self):
self.global_config = GlobalConfig()
self.log = Logger().get('database.database.Database')
def create_default_database(self):
"""
Calls methods needed to create the database.
"""
self.create_db_dir()
self.create_db()
# Execute scripts BEFORE updating schema
run_db_scripts(self.global_config)
self.update_schema()
def create_db_dir(self):
"""
Creates the database directory if it doesn't already exist.
"""
# if database directory does not exist create it
db_path = self.global_config['Database']['path']
(db_dir, db_name) = ntpath.split(db_path)
if not os.path.isdir(db_dir):
self.log.info("Database directory not found, "
"creating database directory...")
os.mkdir(db_dir)
def create_db(self):
"""
Creates the database if it doesn't already exist.
"""
# if database file does not exist in the directory, create it
(db_dir, db_name) = ntpath.split(self.global_config['Database']['path'])
if not os.path.exists(self.global_config['Database']['path']):
self.log.info("Database file not found, creating database file...")
# this actually creates the database file
connection = sqlite3.connect(self.global_config['Database']['path'])
connection.close()
def update_schema(self):
"""
Updates the database when columns have been added to, or
removed from, the schema.
"""
# Create any new tables that have been added to the plugin
# config schema.
db_tables = DataValidator().get_tables()
cfg_tables = get_config_table_list(
self.global_config.get_ports(),
self.global_config.get_plugin_dictionary())
table_diff = list(set(cfg_tables) - set(db_tables))
self.create_non_exist_tables(table_diff)
# Populate the newly created tables with their column
# definitions.
DataValidator().update_tables_and_schema()
self.update_table_structure()
def create_non_exist_tables(self, table_diff):
"""
create tables that do not exist from the table difference between the current database and the configuration
"""
if len(table_diff) > 0:
for table in table_diff:
Table_Init.create_table(table, self.global_config)
self.log.info('Updated database schema, table names now match configuration.')
else:
self.log.info('Database Schema and Configuration table names already match.')
def create_dict_config_column_list(self):
"""
get a dictionary of tables and corresponding columns from the config
"""
config_column_lists = {}
for port in self.global_config.get_ports():
value = self.global_config.get_plugin_dictionary().get(port)
config_column_lists[value.get('table')] = value.get('tableColumns')
return config_column_lists
def create_dict_transformed_column_list(self, database_column_lists):
"""
returns only custom plugin defined columns from database schema i.e.
ignores default columns
"""
transformed_db_column_list = {}
for table in database_column_lists:
col_list = database_column_lists[table]
transformed_db_column_list[table] = []
# default column ids to ignore
default_list = []
for default in default_columns:
default_list.append(default[0])
for column in col_list:
# ignores the default columns
if column[1] in default_list:
continue
transformed_db_column_list[table].append([column[0],column[1],column[2]])
return transformed_db_column_list
def update_table_structure(self):
cfg_schema = self.create_dict_config_column_list()
db_schema = DataValidator().get_schema()
db_schema_sans_defaults = self.create_dict_transformed_column_list(db_schema)
for table in cfg_schema:
if not [(x[1], x[2]) for x in cfg_schema[table]] == \
[(x[1], x[2]) for x in db_schema_sans_defaults[table]]:
Table_Init.change_table_structure(
table, cfg_schema[table], db_schema[table],
self.global_config)
class database_test(unittest.TestCase):
def setUp(self):
self.test_db_dir = '/tests/database/test_database'
self.test_db_file = '/tests/database/test_database/honeyDB.sqlite'
# test configuration files
self.plugins_config_file = 'tests/database/test_config/plugins.cfg'
self.plugins_config_diff_file = 'tests/database/test_config/plugins_diff.cfg'
self.plugins_config_diff_table_file = 'tests/database/test_config/plugins_diff_table.cfg'
self.global_config_file = 'tests/database/test_config/global.cfg'
# create global config instance
self.gci = GlobalConfig(self.plugins_config_file,self.global_config_file)
self.gci.read_global_config()
self.gci.read_plugin_config()
# patch the Logger new method so that it doesn't create the log file, we do not need to test that the logging
# works just that the log statements are called.
@patch.object(Logger,'__new__')
def test_database_init(self,log):
db = database.Database()
self.assertIsInstance(db.global_config,GlobalConfig._GlobalConfig)
self.assertTrue(log.called)
@patch.object(Logger,'__new__')
def test_database_create_default_database(self,log):
db = database.Database()
db.create_default_database()
validator = datavalidator.DataValidator()
# check that the directory exists
self.assertTrue(os.path.isdir(os.getcwd() + self.test_db_dir))
# check that the database file exists
self.assertTrue(os.path.isfile(os.getcwd() + self.test_db_file))
# get the table names from the database
schema_table_list = validator.get_tables()
# get the user defined tables from the configuration file
config_table_list = util.get_config_table_list(self.gci.get_ports(),
self.gci.get_plugin_dictionary())
# check that the non user defined table p0f exists
self.assertTrue('p0f' in schema_table_list)
# check that the non user defined table ipInfo exists
self.assertTrue('ipInfo' in schema_table_list)
# check that the non user defined table sessions exists
self.assertTrue('sessions' in schema_table_list)
# check that the user defined tables are a subset of the tables in the database schema
self.assertTrue(set(config_table_list) < set(schema_table_list))
shutil.rmtree(os.getcwd() + self.test_db_dir)
@patch.object(Logger,'__new__')
def test_database_create_db_dir(self,log):
db = database.Database()
db.create_db_dir()
self.assertTrue(os.path.isdir(os.getcwd() + self.test_db_dir))
self.assertTrue(log.called)
shutil.rmtree(os.getcwd() + self.test_db_dir)
@patch.object(Logger,'__new__')
def test_database_create_db_dir_already_exists(self,log):
os.mkdir(os.getcwd() + self.test_db_dir)
self.assertTrue(os.path.isdir(os.getcwd() + self.test_db_dir))
db = database.Database()
log.reset_mock()
db.create_db_dir()
self.assertFalse(log.called)
shutil.rmtree(os.getcwd() + self.test_db_dir)
@patch.object(Logger,'__new__')
def test_database_create_db(self,log):
db = database.Database()
db.create_db_dir()
db.create_db()
self.assertTrue(os.path.isfile(os.getcwd() + self.test_db_file))
self.assertTrue(log.called)
shutil.rmtree(os.getcwd() + self.test_db_dir)
@patch.object(Logger,'__new__')
def test_database_update_schema(self,log):
db = database.Database()
db.create_db_dir()
db.create_db()
util.run_db_scripts(self.gci)
db.update_schema()
validator = datavalidator.DataValidator()
schema = validator.get_schema()
self.assertTrue(schema['test_http'][5][1] == 'command')
self.assertTrue(schema['test_http2'][6][1] == 'path')
self.assertTrue(len(schema['test_telnet']) == 7)
# set global config instance to the differing column config file
self.gci = GlobalConfig(self.plugins_config_diff_file,self.global_config_file,True)
self.gci.read_global_config()
self.gci.read_plugin_config()
db2 = database.Database()
db2.update_schema()
validator2 = datavalidator.DataValidator()
schema2 = validator2.get_schema()
self.assertTrue(schema2['test_http'][5][1] == 'unit_test_data_1')
self.assertTrue(schema2['test_http2'][6][1] == 'unit_test_data_2')
self.assertTrue(len(schema2['test_telnet']) == 8)
self.assertTrue(schema2['test_telnet'][7][1] == 'unit_test_data_3')
# set global config instance back to normal
self.gci = GlobalConfig(self.plugins_config_file,self.global_config_file,True)
self.gci.read_global_config()
self.gci.read_plugin_config()
shutil.rmtree(os.getcwd() + self.test_db_dir)
