File size: 3,866 Bytes
c7dfe8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import threading
from sqlalchemy import BigInteger, Column, Integer, String, UnicodeText
from Database.sql import BASE, SESSION
DEF_COUNT = 1
DEF_LIMIT = 0
DEF_OBJ = (None, DEF_COUNT, DEF_LIMIT)
class FloodControl(BASE):
__tablename__ = "antiflood"
chat_id = Column(String(14), primary_key=True)
user_id = Column(BigInteger)
count = Column(Integer, default=DEF_COUNT)
limit = Column(Integer, default=DEF_LIMIT)
def __init__(self, chat_id):
self.chat_id = str(chat_id) # ensure string
def __repr__(self):
return "<flood control for %s>" % self.chat_id
class FloodSettings(BASE):
__tablename__ = "antiflood_settings"
chat_id = Column(String(14), primary_key=True)
flood_type = Column(Integer, default=1)
value = Column(UnicodeText, default="0")
def __init__(self, chat_id, flood_type=1, value="0"):
self.chat_id = str(chat_id)
self.flood_type = flood_type
self.value = value
def __repr__(self):
return "<{} will executing {} for flood.>".format(self.chat_id, self.flood_type)
FloodControl.__table__.create(checkfirst=True)
FloodSettings.__table__.create(checkfirst=True)
INSERTION_FLOOD_LOCK = threading.RLock()
INSERTION_FLOOD_SETTINGS_LOCK = threading.RLock()
CHAT_FLOOD = {}
def set_flood(chat_id, amount):
with INSERTION_FLOOD_LOCK:
flood = SESSION.query(FloodControl).get(str(chat_id))
if not flood:
flood = FloodControl(str(chat_id))
flood.user_id = None
flood.limit = amount
CHAT_FLOOD[str(chat_id)] = (None, DEF_COUNT, amount)
SESSION.add(flood)
SESSION.commit()
def update_flood(chat_id: str, user_id) -> bool:
if str(chat_id) in CHAT_FLOOD:
curr_user_id, count, limit = CHAT_FLOOD.get(str(chat_id), DEF_OBJ)
if limit == 0: # no antiflood
return False
if user_id != curr_user_id or user_id is None: # other user
CHAT_FLOOD[str(chat_id)] = (user_id, DEF_COUNT, limit)
return False
count += 1
if count > limit: # too many msgs, kick
CHAT_FLOOD[str(chat_id)] = (None, DEF_COUNT, limit)
return True
# default -> update
CHAT_FLOOD[str(chat_id)] = (user_id, count, limit)
return False
def get_flood_limit(chat_id):
return CHAT_FLOOD.get(str(chat_id), DEF_OBJ)[2]
def set_flood_strength(chat_id, flood_type, value):
# for flood_type
# 1 = ban
# 2 = kick
# 3 = mute
# 4 = tban
# 5 = tmute
with INSERTION_FLOOD_SETTINGS_LOCK:
curr_setting = SESSION.query(FloodSettings).get(str(chat_id))
if not curr_setting:
curr_setting = FloodSettings(
chat_id,
flood_type=int(flood_type),
value=value,
)
curr_setting.flood_type = int(flood_type)
curr_setting.value = str(value)
SESSION.add(curr_setting)
SESSION.commit()
def get_flood_setting(chat_id):
try:
setting = SESSION.query(FloodSettings).get(str(chat_id))
if setting:
return setting.flood_type, setting.value
else:
return 1, "0"
finally:
SESSION.close()
def migrate_chat(old_chat_id, new_chat_id):
with INSERTION_FLOOD_LOCK:
flood = SESSION.query(FloodControl).get(str(old_chat_id))
if flood:
CHAT_FLOOD[str(new_chat_id)] = CHAT_FLOOD.get(str(old_chat_id), DEF_OBJ)
flood.chat_id = str(new_chat_id)
SESSION.commit()
SESSION.close()
def __load_flood_settings():
global CHAT_FLOOD
try:
all_chats = SESSION.query(FloodControl).all()
CHAT_FLOOD = {chat.chat_id: (None, DEF_COUNT, chat.limit) for chat in all_chats}
finally:
SESSION.close()
__load_flood_settings()
|