File size: 2,139 Bytes
44aada1 | 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 | from sqlalchemy import Boolean, Column, String
from telebot.plugins.sql_helper import SESSION, BASE
class Locks(BASE):
__tablename__ = "locks"
chat_id = Column(String(14), primary_key=True)
# Booleans are for "is this locked", _NOT_ "is this allowed"
bots = Column(Boolean, default=False)
commands = Column(Boolean, default=False)
email = Column(Boolean, default=False)
forward = Column(Boolean, default=False)
url = Column(Boolean, default=False)
def __init__(self, chat_id):
self.chat_id = str(chat_id) # ensure string
self.bots = False
self.commands = False
self.email = False
self.forward = False
self.url = False
Locks.__table__.create(checkfirst=True)
def init_locks(chat_id, reset=False):
curr_restr = SESSION.query(Locks).get(str(chat_id))
if reset:
SESSION.delete(curr_restr)
SESSION.flush()
restr = Locks(str(chat_id))
SESSION.add(restr)
SESSION.commit()
return restr
def update_lock(chat_id, lock_type, locked):
curr_perm = SESSION.query(Locks).get(str(chat_id))
if not curr_perm:
curr_perm = init_locks(chat_id)
if lock_type == "bots":
curr_perm.bots = locked
elif lock_type == "commands":
curr_perm.commands = locked
elif lock_type == "email":
curr_perm.email = locked
elif lock_type == "forward":
curr_perm.forward = locked
elif lock_type == "url":
curr_perm.url = locked
SESSION.add(curr_perm)
SESSION.commit()
def is_locked(chat_id, lock_type):
curr_perm = SESSION.query(Locks).get(str(chat_id))
SESSION.close()
if not curr_perm:
return False
elif lock_type == "bots":
return curr_perm.bots
elif lock_type == "commands":
return curr_perm.commands
elif lock_type == "email":
return curr_perm.email
elif lock_type == "forward":
return curr_perm.forward
elif lock_type == "url":
return curr_perm.url
def get_locks(chat_id):
try:
return SESSION.query(Locks).get(str(chat_id))
finally:
SESSION.close()
|