2021-07-27 17:42:47 +00:00
|
|
|
import datetime
|
2022-06-12 14:30:02 +00:00
|
|
|
from typing import Union
|
2021-07-27 17:42:47 +00:00
|
|
|
|
2021-11-12 14:25:53 +00:00
|
|
|
from tenacity import retry, stop_after_attempt
|
|
|
|
|
|
|
|
from config import Config
|
2022-06-12 14:30:02 +00:00
|
|
|
from core.elements.message import MessageSession, FetchTarget
|
2021-08-25 11:23:59 +00:00
|
|
|
from core.elements.temp import EnabledModulesCache, SenderInfoCache
|
2022-06-13 04:49:22 +00:00
|
|
|
from database.orm import Session
|
2022-06-12 14:30:02 +00:00
|
|
|
from database.tables import EnabledModules, MuteList, SenderInfo, TargetAdmin, CommandTriggerTime, GroupAllowList, StoredData
|
2021-09-10 18:05:27 +00:00
|
|
|
|
|
|
|
cache = Config('db_cache')
|
2021-07-11 14:20:24 +00:00
|
|
|
|
|
|
|
|
|
|
|
def convert_list_to_str(lst: list) -> str:
|
2021-09-03 19:22:16 +00:00
|
|
|
filter_lst = []
|
|
|
|
for x in lst:
|
|
|
|
if x != '':
|
|
|
|
filter_lst.append(x)
|
|
|
|
return '|'.join(filter_lst)
|
2021-07-11 14:20:24 +00:00
|
|
|
|
|
|
|
|
|
|
|
def convert_str_to_list(s: str) -> list:
|
|
|
|
return s.split('|')
|
2021-07-11 02:19:40 +00:00
|
|
|
|
|
|
|
|
2021-08-25 11:23:59 +00:00
|
|
|
class Dict2Object(dict):
|
|
|
|
def __getattr__(self, key):
|
|
|
|
return self.get(key)
|
|
|
|
|
|
|
|
def __setattr__(self, key, value):
|
|
|
|
self[key] = value
|
|
|
|
|
|
|
|
|
2022-06-13 04:49:22 +00:00
|
|
|
session = Session.session
|
2021-08-24 16:49:19 +00:00
|
|
|
|
|
|
|
|
2021-11-03 11:27:21 +00:00
|
|
|
def auto_rollback_error(func):
|
|
|
|
def wrapper(*args, **kwargs):
|
|
|
|
try:
|
|
|
|
return func(*args, **kwargs)
|
|
|
|
except Exception as e:
|
|
|
|
session.rollback()
|
|
|
|
raise e
|
2021-11-12 14:25:53 +00:00
|
|
|
|
2021-11-03 11:27:21 +00:00
|
|
|
return wrapper
|
|
|
|
|
|
|
|
|
2021-07-11 02:19:40 +00:00
|
|
|
class BotDBUtil:
|
|
|
|
class Module:
|
2021-09-13 04:39:49 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-08-21 15:58:07 +00:00
|
|
|
def __init__(self, msg: [MessageSession, str]):
|
|
|
|
if isinstance(msg, MessageSession):
|
|
|
|
self.targetId = str(msg.target.targetId)
|
|
|
|
else:
|
|
|
|
self.targetId = msg
|
2021-08-25 11:23:59 +00:00
|
|
|
self.need_insert = False
|
2021-09-10 18:05:27 +00:00
|
|
|
self.enable_modules_list = EnabledModulesCache.get_cache(self.targetId) if cache else False
|
2021-08-25 11:23:59 +00:00
|
|
|
if not self.enable_modules_list:
|
|
|
|
query = self.query_EnabledModules
|
|
|
|
if query is None:
|
|
|
|
self.need_insert = True
|
|
|
|
self.enable_modules_list = []
|
|
|
|
else:
|
|
|
|
query_ = query.enabledModules
|
|
|
|
self.enable_modules_list = convert_str_to_list(query_)
|
2021-09-10 18:05:27 +00:00
|
|
|
if cache:
|
|
|
|
EnabledModulesCache.add_cache(self.targetId, self.enable_modules_list)
|
2021-08-25 11:23:59 +00:00
|
|
|
|
|
|
|
@property
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-08-25 11:23:59 +00:00
|
|
|
def query_EnabledModules(self):
|
2021-11-03 11:27:21 +00:00
|
|
|
return session.query(EnabledModules).filter_by(targetId=self.targetId).first()
|
2021-07-11 02:19:40 +00:00
|
|
|
|
2021-08-25 11:45:03 +00:00
|
|
|
def check_target_enabled_module_list(self) -> list:
|
|
|
|
return self.enable_modules_list
|
|
|
|
|
2021-07-11 02:19:40 +00:00
|
|
|
def check_target_enabled_module(self, module_name) -> bool:
|
|
|
|
return True if module_name in self.enable_modules_list else False
|
|
|
|
|
2021-09-12 12:51:13 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-07-11 02:19:40 +00:00
|
|
|
def enable(self, module_name) -> bool:
|
2021-11-03 11:27:21 +00:00
|
|
|
if isinstance(module_name, str):
|
|
|
|
if module_name not in self.enable_modules_list:
|
|
|
|
self.enable_modules_list.append(module_name)
|
|
|
|
elif isinstance(module_name, (list, tuple)):
|
|
|
|
for x in module_name:
|
|
|
|
if x not in self.enable_modules_list:
|
|
|
|
self.enable_modules_list.append(x)
|
|
|
|
value = convert_list_to_str(self.enable_modules_list)
|
|
|
|
if self.need_insert:
|
|
|
|
table = EnabledModules(targetId=self.targetId,
|
|
|
|
enabledModules=value)
|
|
|
|
session.add_all([table])
|
|
|
|
else:
|
|
|
|
self.query_EnabledModules.enabledModules = value
|
|
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
|
|
if cache:
|
|
|
|
EnabledModulesCache.add_cache(self.targetId, self.enable_modules_list)
|
|
|
|
return True
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def disable(self, module_name) -> bool:
|
|
|
|
if isinstance(module_name, str):
|
|
|
|
if module_name in self.enable_modules_list:
|
|
|
|
self.enable_modules_list.remove(module_name)
|
|
|
|
elif isinstance(module_name, (list, tuple)):
|
|
|
|
for x in module_name:
|
|
|
|
if x in self.enable_modules_list:
|
|
|
|
self.enable_modules_list.remove(x)
|
|
|
|
if not self.need_insert:
|
|
|
|
self.query_EnabledModules.enabledModules = convert_list_to_str(self.enable_modules_list)
|
2021-07-11 02:19:40 +00:00
|
|
|
session.commit()
|
2021-08-24 10:22:36 +00:00
|
|
|
session.expire_all()
|
2021-09-10 18:05:27 +00:00
|
|
|
if cache:
|
|
|
|
EnabledModulesCache.add_cache(self.targetId, self.enable_modules_list)
|
2021-11-03 11:27:21 +00:00
|
|
|
return True
|
2021-07-12 13:31:11 +00:00
|
|
|
|
2021-08-02 13:39:50 +00:00
|
|
|
@staticmethod
|
2021-11-18 10:24:15 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
2021-08-02 13:39:50 +00:00
|
|
|
def get_enabled_this(module_name):
|
|
|
|
query = session.query(EnabledModules).filter(EnabledModules.enabledModules.like(f'%{module_name}%'))
|
|
|
|
targetIds = []
|
|
|
|
for x in query:
|
|
|
|
enabled_list = convert_str_to_list(x.enabledModules)
|
|
|
|
if module_name in enabled_list:
|
|
|
|
targetIds.append(x.targetId)
|
|
|
|
return targetIds
|
|
|
|
|
2021-07-15 14:59:32 +00:00
|
|
|
class SenderInfo:
|
2021-09-12 12:51:13 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-07-16 15:50:25 +00:00
|
|
|
def __init__(self, senderId):
|
|
|
|
self.senderId = senderId
|
2021-09-10 18:05:27 +00:00
|
|
|
query_cache = SenderInfoCache.get_cache(self.senderId) if cache else False
|
2021-08-25 11:23:59 +00:00
|
|
|
if query_cache:
|
|
|
|
self.query = Dict2Object(query_cache)
|
|
|
|
else:
|
|
|
|
self.query = self.query_SenderInfo
|
2021-11-03 11:27:21 +00:00
|
|
|
if self.query is None:
|
|
|
|
session.add_all([SenderInfo(id=senderId)])
|
|
|
|
session.commit()
|
|
|
|
self.query = session.query(SenderInfo).filter_by(id=senderId).first()
|
|
|
|
if cache:
|
|
|
|
SenderInfoCache.add_cache(self.senderId, self.query.__dict__)
|
2021-08-25 11:23:59 +00:00
|
|
|
|
|
|
|
@property
|
2021-11-18 10:24:57 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-08-25 11:23:59 +00:00
|
|
|
def query_SenderInfo(self):
|
2021-11-03 11:27:21 +00:00
|
|
|
return session.query(SenderInfo).filter_by(id=self.senderId).first()
|
2021-07-13 15:48:43 +00:00
|
|
|
|
2021-09-12 12:51:13 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-07-16 15:50:25 +00:00
|
|
|
def edit(self, column: str, value):
|
2021-11-03 11:27:21 +00:00
|
|
|
query = self.query_SenderInfo
|
|
|
|
setattr(query, column, value)
|
|
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
|
|
if cache:
|
|
|
|
SenderInfoCache.add_cache(self.senderId, query.__dict__)
|
|
|
|
return True
|
2021-07-13 15:48:43 +00:00
|
|
|
|
2021-11-03 11:27:21 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
2021-07-24 08:59:15 +00:00
|
|
|
def check_TargetAdmin(self, targetId):
|
|
|
|
query = session.query(TargetAdmin).filter_by(senderId=self.senderId, targetId=targetId).first()
|
|
|
|
if query is not None:
|
|
|
|
return query
|
|
|
|
return False
|
|
|
|
|
2021-09-12 12:51:13 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-07-24 08:59:15 +00:00
|
|
|
def add_TargetAdmin(self, targetId):
|
2021-11-03 11:27:21 +00:00
|
|
|
if not self.check_TargetAdmin(targetId):
|
|
|
|
session.add_all([TargetAdmin(senderId=self.senderId, targetId=targetId)])
|
|
|
|
session.commit()
|
|
|
|
return True
|
2021-07-24 08:59:15 +00:00
|
|
|
|
2021-09-12 12:51:13 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-07-24 08:59:15 +00:00
|
|
|
def remove_TargetAdmin(self, targetId):
|
2021-11-03 11:27:21 +00:00
|
|
|
query = self.check_TargetAdmin(targetId)
|
|
|
|
if query:
|
|
|
|
session.delete(query)
|
|
|
|
session.commit()
|
2022-01-30 14:02:17 +00:00
|
|
|
return True
|
2021-08-26 16:24:21 +00:00
|
|
|
|
2021-07-27 17:42:47 +00:00
|
|
|
class CoolDown:
|
2021-11-03 11:27:21 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
2021-07-27 17:42:47 +00:00
|
|
|
def __init__(self, msg: MessageSession, name):
|
|
|
|
self.msg = msg
|
|
|
|
self.name = name
|
2022-05-18 11:47:35 +00:00
|
|
|
self.query = session.query(CommandTriggerTime).filter_by(targetId=str(msg.target.senderId),
|
2021-08-07 07:56:48 +00:00
|
|
|
commandName=name).first()
|
2021-07-27 17:42:47 +00:00
|
|
|
self.need_insert = True if self.query is None else False
|
|
|
|
|
|
|
|
def check(self, delay):
|
|
|
|
if not self.need_insert:
|
|
|
|
now = datetime.datetime.now().timestamp() - self.query.timestamp.timestamp()
|
|
|
|
if now > delay:
|
|
|
|
return 0
|
|
|
|
return now
|
|
|
|
return 0
|
|
|
|
|
2021-09-12 12:51:13 +00:00
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-07-27 17:42:47 +00:00
|
|
|
def reset(self):
|
2021-11-03 11:27:21 +00:00
|
|
|
if not self.need_insert:
|
|
|
|
session.delete(self.query)
|
2021-07-27 17:42:47 +00:00
|
|
|
session.commit()
|
2022-05-18 11:52:16 +00:00
|
|
|
session.add_all([CommandTriggerTime(targetId=self.msg.target.senderId, commandName=self.name)])
|
2021-11-03 11:27:21 +00:00
|
|
|
session.commit()
|
2021-10-16 16:17:21 +00:00
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
2021-11-03 11:27:21 +00:00
|
|
|
@auto_rollback_error
|
2021-11-15 16:26:11 +00:00
|
|
|
def isGroupInAllowList(targetId):
|
2021-11-03 11:27:21 +00:00
|
|
|
session.expire_all()
|
2021-11-15 16:26:11 +00:00
|
|
|
query = session.query(GroupAllowList).filter_by(targetId=targetId).first()
|
2021-11-03 11:27:21 +00:00
|
|
|
if query is not None:
|
|
|
|
return True
|
|
|
|
return False
|
2021-10-16 16:17:21 +00:00
|
|
|
|
2022-01-08 11:05:17 +00:00
|
|
|
class Muting:
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def __init__(self, msg: MessageSession):
|
|
|
|
self.msg = msg
|
|
|
|
self.targetId = msg.target.targetId
|
|
|
|
self.query = session.query(MuteList).filter_by(targetId=self.targetId).first()
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def check(self):
|
|
|
|
if self.query is not None:
|
|
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def add(self):
|
|
|
|
session.add(MuteList(targetId=self.targetId))
|
|
|
|
session.commit()
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def remove(self):
|
|
|
|
if self.query is not None:
|
|
|
|
session.delete(self.query)
|
|
|
|
session.commit()
|
|
|
|
|
2022-06-12 14:30:02 +00:00
|
|
|
class Data:
|
|
|
|
def __init__(self, msg: Union[MessageSession, FetchTarget]):
|
|
|
|
self.targetName = msg.target.clientName if isinstance(msg, MessageSession) else msg.name
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def add(self, name, value: str):
|
|
|
|
session.add(StoredData(name=f'{self.targetName}|{name}', value=value))
|
|
|
|
session.commit()
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def get(self, name):
|
|
|
|
return session.query(StoredData).filter_by(name=f'{self.targetName}|{name}').first()
|
|
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3))
|
|
|
|
@auto_rollback_error
|
|
|
|
def update(self, name, value: str):
|
|
|
|
exists = self.get(name)
|
|
|
|
if exists is None:
|
|
|
|
self.add(name=name, value=value)
|
|
|
|
else:
|
|
|
|
exists.value = value
|
|
|
|
session.commit()
|
|
|
|
return True
|
|
|
|
|
2022-01-08 11:05:17 +00:00
|
|
|
|
2021-11-03 11:27:21 +00:00
|
|
|
__all__ = ["BotDBUtil", "auto_rollback_error", "session"]
|