mirror of
https://github.com/iio612/DEFENDER.git
synced 2026-02-13 11:14:23 +00:00
Introduce full asyncio version (still some module to migrate). Defender JSONRPC Server ready and using with uvcorn
This commit is contained in:
103
core/base.py
103
core/base.py
@@ -45,9 +45,6 @@ class Base:
|
||||
# Liste des fonctions en attentes
|
||||
self.periodic_func: dict[object] = self.Settings.PERIODIC_FUNC
|
||||
|
||||
# Création du lock
|
||||
self.lock = self.Settings.LOCK
|
||||
|
||||
# Init install variable
|
||||
self.install: bool = False
|
||||
|
||||
@@ -57,8 +54,8 @@ class Base:
|
||||
# Create the database
|
||||
# self.__create_db()
|
||||
|
||||
def init(self) -> None:
|
||||
self.__create_db()
|
||||
async def init(self) -> None:
|
||||
await self.__create_db()
|
||||
|
||||
def __set_current_defender_version(self) -> None:
|
||||
"""This will put the current version of Defender
|
||||
@@ -145,7 +142,7 @@ class Base:
|
||||
except Exception as err:
|
||||
self.logs.error(f'General Error: {err}')
|
||||
|
||||
def create_log(self, log_message: str) -> None:
|
||||
async def create_log(self, log_message: str) -> None:
|
||||
"""Enregiste les logs
|
||||
|
||||
Args:
|
||||
@@ -156,11 +153,11 @@ class Base:
|
||||
"""
|
||||
sql_insert = f"INSERT INTO {self.Config.TABLE_LOG} (datetime, server_msg) VALUES (:datetime, :server_msg)"
|
||||
mes_donnees = {'datetime': str(self.Utils.get_sdatetime()),'server_msg': f'{log_message}'}
|
||||
self.db_execute_query(sql_insert, mes_donnees)
|
||||
await self.db_execute_query(sql_insert, mes_donnees)
|
||||
|
||||
return None
|
||||
|
||||
def log_cmd(self, user_cmd: str, cmd: str) -> None:
|
||||
async def log_cmd(self, user_cmd: str, cmd: str) -> None:
|
||||
"""Enregistre les commandes envoyées par les utilisateurs
|
||||
|
||||
Args:
|
||||
@@ -176,11 +173,11 @@ class Base:
|
||||
|
||||
insert_cmd_query = f"INSERT INTO {self.Config.TABLE_COMMAND} (datetime, user, commande) VALUES (:datetime, :user, :commande)"
|
||||
mes_donnees = {'datetime': self.Utils.get_sdatetime(), 'user': user_cmd, 'commande': cmd}
|
||||
self.db_execute_query(insert_cmd_query, mes_donnees)
|
||||
await self.db_execute_query(insert_cmd_query, mes_donnees)
|
||||
|
||||
return None
|
||||
|
||||
def db_sync_core_config(self, module_name: str, dataclassObj: object) -> bool:
|
||||
async def db_sync_core_config(self, module_name: str, dataclassObj: object) -> bool:
|
||||
"""Sync module local parameters with the database
|
||||
if new module then local param will be stored in the database
|
||||
if old module then db param will be moved to the local dataclassObj
|
||||
@@ -207,7 +204,7 @@ class Base:
|
||||
param_to_search = {'module_name': module_name, 'param_key': param_key}
|
||||
|
||||
search_query = f'''SELECT id FROM {core_table} WHERE module_name = :module_name AND param_key = :param_key'''
|
||||
excecute_search_query = self.db_execute_query(search_query, param_to_search)
|
||||
excecute_search_query = await self.db_execute_query(search_query, param_to_search)
|
||||
result_search_query = excecute_search_query.fetchone()
|
||||
|
||||
if result_search_query is None:
|
||||
@@ -219,7 +216,7 @@ class Base:
|
||||
insert_query = f'''INSERT INTO {core_table} (datetime, module_name, param_key, param_value)
|
||||
VALUES (:datetime, :module_name, :param_key, :param_value)
|
||||
'''
|
||||
execution = self.db_execute_query(insert_query, param_to_insert)
|
||||
execution = await self.db_execute_query(insert_query, param_to_insert)
|
||||
|
||||
if execution.rowcount > 0:
|
||||
self.logs.debug(f'New parameter added to the database: {param_key} --> {param_value}')
|
||||
@@ -227,14 +224,14 @@ class Base:
|
||||
# Delete from DB unused parameter
|
||||
query_select = f"SELECT module_name, param_key, param_value FROM {core_table} WHERE module_name = :module_name"
|
||||
parameter = {'module_name': module_name}
|
||||
execute_query_select = self.db_execute_query(query_select, parameter)
|
||||
execute_query_select = await self.db_execute_query(query_select, parameter)
|
||||
result_query_select = execute_query_select.fetchall()
|
||||
|
||||
for result in result_query_select:
|
||||
db_mod_name, db_param_key, db_param_value = result
|
||||
if not hasattr(dataclassObj, db_param_key):
|
||||
mes_donnees = {'param_key': db_param_key, 'module_name': db_mod_name}
|
||||
execute_delete = self.db_execute_query(f'DELETE FROM {core_table} WHERE module_name = :module_name and param_key = :param_key', mes_donnees)
|
||||
execute_delete = await self.db_execute_query(f'DELETE FROM {core_table} WHERE module_name = :module_name and param_key = :param_key', mes_donnees)
|
||||
row_affected = execute_delete.rowcount
|
||||
if row_affected > 0:
|
||||
self.logs.debug(f'A parameter has been deleted from the database: {db_param_key} --> {db_param_value} | Mod: {db_mod_name}')
|
||||
@@ -242,7 +239,7 @@ class Base:
|
||||
# Sync local variable with Database
|
||||
query = f"SELECT param_key, param_value FROM {core_table} WHERE module_name = :module_name"
|
||||
parameter = {'module_name': module_name}
|
||||
response = self.db_execute_query(query, parameter)
|
||||
response = await self.db_execute_query(query, parameter)
|
||||
result = response.fetchall()
|
||||
|
||||
for param, value in result:
|
||||
@@ -259,7 +256,7 @@ class Base:
|
||||
self.logs.error(err)
|
||||
return False
|
||||
|
||||
def db_update_core_config(self, module_name:str, dataclass_obj: object, param_key:str, param_value: str) -> bool:
|
||||
async def db_update_core_config(self, module_name:str, dataclass_obj: object, param_key:str, param_value: str) -> bool:
|
||||
|
||||
core_table = self.Config.TABLE_CONFIG
|
||||
# Check if the param exist
|
||||
@@ -269,7 +266,7 @@ class Base:
|
||||
|
||||
mes_donnees = {'module_name': module_name, 'param_key': param_key, 'param_value': param_value}
|
||||
search_param_query = f"SELECT id FROM {core_table} WHERE module_name = :module_name AND param_key = :param_key"
|
||||
result = self.db_execute_query(search_param_query, mes_donnees)
|
||||
result = await self.db_execute_query(search_param_query, mes_donnees)
|
||||
is_param_exist = result.fetchone()
|
||||
|
||||
if not is_param_exist is None:
|
||||
@@ -279,7 +276,7 @@ class Base:
|
||||
'param_value': param_value
|
||||
}
|
||||
query = f'''UPDATE {core_table} SET datetime = :datetime, param_value = :param_value WHERE module_name = :module_name AND param_key = :param_key'''
|
||||
update = self.db_execute_query(query, mes_donnees)
|
||||
update = await self.db_execute_query(query, mes_donnees)
|
||||
updated_rows = update.rowcount
|
||||
if updated_rows > 0:
|
||||
setattr(dataclass_obj, param_key, self.int_if_possible(param_value))
|
||||
@@ -293,9 +290,9 @@ class Base:
|
||||
|
||||
return True
|
||||
|
||||
def db_create_first_admin(self) -> None:
|
||||
async def db_create_first_admin(self) -> None:
|
||||
|
||||
user = self.db_execute_query(f"SELECT id FROM {self.Config.TABLE_ADMIN}")
|
||||
user = await self.db_execute_query(f"SELECT id FROM {self.Config.TABLE_ADMIN}")
|
||||
if not user.fetchall():
|
||||
admin = self.Config.OWNER
|
||||
password = self.Utils.hash_password(self.Config.PASSWORD)
|
||||
@@ -308,7 +305,7 @@ class Base:
|
||||
'language': 'EN',
|
||||
'level': 5
|
||||
}
|
||||
self.db_execute_query(f"""
|
||||
await self.db_execute_query(f"""
|
||||
INSERT INTO {self.Config.TABLE_ADMIN}
|
||||
(createdOn, user, password, hostname, vhost, language, level)
|
||||
VALUES
|
||||
@@ -381,11 +378,28 @@ class Base:
|
||||
return None
|
||||
|
||||
task = asyncio.create_task(func, name=name)
|
||||
task.add_done_callback(self.asynctask_done)
|
||||
self.running_asynctasks.append(task)
|
||||
|
||||
self.logs.debug(f"++ New asynchrone task created as: {task.get_name()}")
|
||||
return task
|
||||
|
||||
def asynctask_done(self, task: asyncio.Task):
|
||||
"""Log task when done
|
||||
|
||||
Args:
|
||||
task (asyncio.Task): The Asyncio Task callback
|
||||
"""
|
||||
try:
|
||||
if task.exception():
|
||||
self.logs.error(f"[ASYNCIO] Task {task.get_name()} failed with exception: {task.exception()}")
|
||||
else:
|
||||
self.logs.debug(f"[ASYNCIO] Task {task.get_name()} completed successfully.")
|
||||
except asyncio.CancelledError as ce:
|
||||
self.logs.debug(f"[ASYNCIO] Task {task.get_name()} terminated with cancelled error.")
|
||||
except asyncio.InvalidStateError as ie:
|
||||
self.logs.debug(f"[ASYNCIO] Task {task.get_name()} terminated with invalid state error.")
|
||||
|
||||
def is_thread_alive(self, thread_name: str) -> bool:
|
||||
"""Check if the thread is still running! using the is_alive method of Threads.
|
||||
|
||||
@@ -429,7 +443,7 @@ class Base:
|
||||
Returns:
|
||||
int: Number of threads
|
||||
"""
|
||||
with self.lock:
|
||||
with self.Settings.LOCK:
|
||||
count = 0
|
||||
|
||||
for thr in self.running_threads:
|
||||
@@ -478,7 +492,7 @@ class Base:
|
||||
self.running_sockets.remove(soc)
|
||||
self.logs.debug(f"-- Socket ==> closed {str(soc.fileno())}")
|
||||
|
||||
def shutdown(self) -> None:
|
||||
async def shutdown(self) -> None:
|
||||
"""Methode qui va préparer l'arrêt complêt du service
|
||||
"""
|
||||
# Nettoyage des timers
|
||||
@@ -507,6 +521,9 @@ class Base:
|
||||
|
||||
self.running_sockets.remove(soc)
|
||||
self.logs.debug(f"> Socket ==> closed {str(soc.fileno())}")
|
||||
|
||||
await self.Loader.RpcServer.stop_server()
|
||||
self.db_close()
|
||||
|
||||
return None
|
||||
|
||||
@@ -524,7 +541,7 @@ class Base:
|
||||
self.logs.info("-- Database connexion has been initiated")
|
||||
return engine, cursor
|
||||
|
||||
def __create_db(self) -> None:
|
||||
async def __create_db(self) -> None:
|
||||
|
||||
table_core_log = f'''CREATE TABLE IF NOT EXISTS {self.Config.TABLE_LOG} (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
@@ -594,27 +611,27 @@ class Base:
|
||||
)
|
||||
'''
|
||||
|
||||
self.db_execute_query(table_core_log)
|
||||
self.db_execute_query(table_core_log_command)
|
||||
self.db_execute_query(table_core_module)
|
||||
self.db_execute_query(table_core_admin)
|
||||
self.db_execute_query(table_core_client)
|
||||
self.db_execute_query(table_core_channel)
|
||||
self.db_execute_query(table_core_config)
|
||||
await self.db_execute_query(table_core_log)
|
||||
await self.db_execute_query(table_core_log_command)
|
||||
await self.db_execute_query(table_core_module)
|
||||
await self.db_execute_query(table_core_admin)
|
||||
await self.db_execute_query(table_core_client)
|
||||
await self.db_execute_query(table_core_channel)
|
||||
await self.db_execute_query(table_core_config)
|
||||
|
||||
# Patch database
|
||||
self.db_patch(self.Config.TABLE_ADMIN, "language", "TEXT")
|
||||
await self.db_patch(self.Config.TABLE_ADMIN, "language", "TEXT")
|
||||
|
||||
if self.install:
|
||||
self.Loader.ModuleUtils.db_register_module('mod_command', 'sys', True)
|
||||
self.Loader.ModuleUtils.db_register_module('mod_defender', 'sys', True)
|
||||
await self.Loader.ModuleUtils.db_register_module('mod_command', 'sys', True)
|
||||
await self.Loader.ModuleUtils.db_register_module('mod_defender', 'sys', True)
|
||||
self.install = False
|
||||
|
||||
return None
|
||||
|
||||
def db_execute_query(self, query:str, params:dict = {}) -> CursorResult:
|
||||
async def db_execute_query(self, query:str, params:dict = {}) -> CursorResult:
|
||||
|
||||
with self.lock:
|
||||
async with self.Loader.Settings.AILOCK:
|
||||
insert_query = text(query)
|
||||
if not params:
|
||||
response = self.cursor.execute(insert_query)
|
||||
@@ -625,8 +642,8 @@ class Base:
|
||||
|
||||
return response
|
||||
|
||||
def db_is_column_exist(self, table_name: str, column_name: str) -> bool:
|
||||
q = self.db_execute_query(f"PRAGMA table_info({table_name})")
|
||||
async def db_is_column_exist(self, table_name: str, column_name: str) -> bool:
|
||||
q = await self.db_execute_query(f"PRAGMA table_info({table_name})")
|
||||
existing_columns = [col[1] for col in q.fetchall()]
|
||||
|
||||
if column_name in existing_columns:
|
||||
@@ -634,12 +651,12 @@ class Base:
|
||||
else:
|
||||
return False
|
||||
|
||||
def db_patch(self, table_name: str, column_name: str, column_type: str) -> bool:
|
||||
if not self.db_is_column_exist(table_name, column_name):
|
||||
async def db_patch(self, table_name: str, column_name: str, column_type: str) -> bool:
|
||||
if not await self.db_is_column_exist(table_name, column_name):
|
||||
patch = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
|
||||
update_row = f"UPDATE {table_name} SET language = 'EN' WHERE language is null"
|
||||
self.db_execute_query(patch)
|
||||
self.db_execute_query(update_row)
|
||||
await self.db_execute_query(patch)
|
||||
await self.db_execute_query(update_row)
|
||||
self.logs.debug(f"The patch has been applied")
|
||||
self.logs.debug(f"Table name: {table_name}, Column name: {column_name}, Column type: {column_type}")
|
||||
return True
|
||||
@@ -647,9 +664,9 @@ class Base:
|
||||
return False
|
||||
|
||||
def db_close(self) -> None:
|
||||
|
||||
try:
|
||||
self.cursor.close()
|
||||
self.logs.debug("Database engine closed!")
|
||||
except AttributeError as ae:
|
||||
self.logs.error(f"Attribute Error : {ae}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user