437 lines
14 KiB
Python
437 lines
14 KiB
Python
# Telebot imports
|
|
from telebot.async_telebot import AsyncTeleBot
|
|
from telebot.asyncio_storage import StateMemoryStorage, StatePickleStorage
|
|
from telebot.asyncio_handler_backends import State, StatesGroup
|
|
from telebot.asyncio_filters import StateFilter
|
|
from telebot import types
|
|
from telebot.callback_data import CallbackData, CallbackDataFilter
|
|
from telebot.types import User as TgUser
|
|
from telebot.types import Chat as TgChat
|
|
from telebot.types import Message
|
|
from telebot.util import user_link
|
|
|
|
# Async things imports
|
|
import asyncio
|
|
|
|
# Other modules imports
|
|
from datetime import datetime
|
|
import math
|
|
from typing import Union
|
|
|
|
# Local imports
|
|
from config import token
|
|
import textbook
|
|
import keyboards
|
|
|
|
# DB
|
|
from db.base import Session, engine, Base
|
|
from db.models import User, Group, GroupMember, Fund, FundMember
|
|
|
|
|
|
bot = AsyncTeleBot(token, state_storage=StatePickleStorage())
|
|
session: Session = None
|
|
|
|
|
|
class States(StatesGroup):
|
|
default = State()
|
|
newfund_amount = State()
|
|
close_fund = State()
|
|
|
|
|
|
# Utils
|
|
|
|
|
|
def get_fund_text(fund: Fund):
|
|
count = fund.members.count()
|
|
contributors = fund.members.filter(FundMember.contributed).count()
|
|
personal_amount = math.ceil(fund.amount / count)
|
|
return textbook.fund.format(
|
|
active="🟢" if fund.active else "🔴",
|
|
name=fund.name,
|
|
amount=fund.amount,
|
|
personal_amount=personal_amount,
|
|
collected_amount=personal_amount * contributors,
|
|
contributors=contributors,
|
|
count=count,
|
|
)
|
|
|
|
|
|
def get_user(tg_user: TgUser) -> User:
|
|
if user := session.query(User).filter(User.id == tg_user.id).first():
|
|
return user
|
|
user = User(
|
|
id=tg_user.id,
|
|
name=tg_user.first_name,
|
|
username=tg_user.username,
|
|
)
|
|
session.add(user)
|
|
session.commit()
|
|
return user
|
|
|
|
|
|
def get_group(tg_chat: TgChat) -> Group:
|
|
if group := session.query(Group).filter(Group.id == tg_chat.id).first():
|
|
return group
|
|
group = Group(
|
|
id=tg_chat.id,
|
|
)
|
|
session.add(group)
|
|
session.commit()
|
|
return group
|
|
|
|
|
|
# Bot logic
|
|
|
|
|
|
@bot.message_handler(commands=["start"])
|
|
async def start(msg: Message):
|
|
if msg.chat.type == "private":
|
|
user = get_user(msg.from_user)
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.start_private.format(count=user.fund_members.count()),
|
|
)
|
|
|
|
elif msg.chat.type in ("group", "supergroup"):
|
|
get_group(msg.chat)
|
|
await bot.send_message(chat_id=msg.chat.id, text=textbook.start_group)
|
|
|
|
|
|
@bot.message_handler(commands=["setup"])
|
|
async def setup(msg: Message):
|
|
get_group(msg.chat)
|
|
if msg.chat.type in ("group", "supergroup"):
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id, text=textbook.setup, reply_markup=keyboards.setup()
|
|
)
|
|
|
|
|
|
@bot.callback_query_handler(func=lambda c: c.data == "register_group_member")
|
|
async def register_group_member(call: types.CallbackQuery):
|
|
new = False
|
|
get_user(call.from_user)
|
|
if (
|
|
group_member := session.query(GroupMember)
|
|
.filter(
|
|
GroupMember.group_id == call.message.chat.id,
|
|
GroupMember.user_id == call.from_user.id,
|
|
)
|
|
.first()
|
|
):
|
|
session.delete(group_member)
|
|
else:
|
|
group_member = GroupMember(
|
|
group_id=call.message.chat.id,
|
|
user_id=call.from_user.id,
|
|
)
|
|
session.add(group_member)
|
|
new = True
|
|
session.commit()
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.user_parted if new else textbook.user_left,
|
|
show_alert=True,
|
|
)
|
|
|
|
|
|
@bot.message_handler(commands=["newfund"])
|
|
async def newfund(msg: Message):
|
|
if msg.chat.type in ("group", "supergroup"):
|
|
if group := get_group(msg.chat):
|
|
if group.group_members.count():
|
|
if (
|
|
session.query(Fund)
|
|
.filter(Fund.group_id == msg.chat.id, Fund.active)
|
|
.first()
|
|
):
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.newfund_already_exists,
|
|
)
|
|
else:
|
|
await bot.set_state(
|
|
user_id=msg.from_user.id,
|
|
chat_id=msg.chat.id,
|
|
state=States.newfund_amount,
|
|
)
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.newfund_amount.format(
|
|
user=user_link(msg.from_user)
|
|
),
|
|
reply_markup=keyboards.cancel(),
|
|
parse_mode="html",
|
|
)
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.not_set_up,
|
|
)
|
|
|
|
|
|
@bot.callback_query_handler(
|
|
func=lambda c: c.data == "cancel", state=States.newfund_amount
|
|
)
|
|
async def cancel_newfund_amount(call: types.CallbackQuery):
|
|
await bot.set_state(
|
|
user_id=call.from_user.id, chat_id=call.message.chat.id, state=States.default
|
|
)
|
|
await bot.delete_message(chat_id=call.message.chat.id, message_id=call.message.id)
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.cancel,
|
|
)
|
|
|
|
|
|
@bot.message_handler(content_types=["text"], state=States.newfund_amount)
|
|
async def newfund_amount(msg: Message):
|
|
if not session.query(User).filter(User.id == msg.from_user.id).first():
|
|
user = User(
|
|
id=msg.from_user.id,
|
|
name=msg.from_user.first_name,
|
|
username=msg.from_user.username,
|
|
)
|
|
session.add(user)
|
|
session.commit()
|
|
if msg.text.isdigit():
|
|
await bot.set_state(
|
|
user_id=msg.from_user.id, chat_id=msg.chat.id, state=States.default
|
|
)
|
|
fund = Fund(
|
|
owner_id=msg.from_user.id, group_id=msg.chat.id, amount=int(msg.text)
|
|
)
|
|
session.add(fund)
|
|
for group_member in session.query(GroupMember).filter(
|
|
GroupMember.group_id == msg.chat.id
|
|
):
|
|
fund_member = FundMember(user_id=group_member.user.id, fund_id=fund.id)
|
|
session.add(fund_member)
|
|
session.commit()
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.fund_created.format(fund=fund.name),
|
|
)
|
|
await asyncio.sleep(1)
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=get_fund_text(fund),
|
|
reply_markup=keyboards.fund_markup(),
|
|
parse_mode="html",
|
|
)
|
|
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.not_number.format(user=user_link(msg.from_user)),
|
|
reply_markup=keyboards.cancel(),
|
|
parse_mode="html",
|
|
)
|
|
|
|
|
|
@bot.message_handler(commands=["fund"])
|
|
async def fund(msg: Message):
|
|
if msg.chat.type in ("group", "supergroup"):
|
|
group = get_group(msg.chat)
|
|
if (
|
|
fund := session.query(Fund)
|
|
.filter(Fund.group_id == group.id, Fund.active == True)
|
|
.first()
|
|
):
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=get_fund_text(fund),
|
|
reply_markup=keyboards.fund_markup(),
|
|
parse_mode="html",
|
|
)
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.fund_not_found,
|
|
)
|
|
|
|
|
|
@bot.callback_query_handler(func=lambda c: c.data == "close_fund")
|
|
async def close_fund_prompt(call: types.CallbackQuery):
|
|
if group := get_group(call.message.chat):
|
|
if fund := group.funds.filter(Fund.active).first():
|
|
if call.from_user.id == fund.owner_id:
|
|
await bot.edit_message_text(
|
|
text=textbook.close_fund_prompt,
|
|
chat_id=call.message.chat.id,
|
|
message_id=call.message.id,
|
|
reply_markup=keyboards.yes_no(),
|
|
parse_mode="html",
|
|
)
|
|
await bot.set_state(
|
|
user_id=call.from_user.id,
|
|
chat_id=call.message.chat.id,
|
|
state=States.close_fund,
|
|
)
|
|
else:
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.not_owner.format(owner_name=fund.owner.name),
|
|
)
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=call.message.chat.id,
|
|
text=textbook.fund_not_found,
|
|
)
|
|
|
|
|
|
@bot.callback_query_handler(
|
|
func=lambda c: c.data in ("yes", "no"), state=States.close_fund
|
|
)
|
|
async def close_fund(call: types.CallbackQuery):
|
|
if group := get_group(call.message.chat):
|
|
if fund := group.funds.filter(Fund.active).first():
|
|
if call.from_user.id == fund.owner_id:
|
|
if call.data == "yes":
|
|
setattr(fund, "active", False)
|
|
session.commit()
|
|
await bot.edit_message_text(
|
|
text=textbook.fund_closed.format(
|
|
name=fund.name, fund_text=get_fund_text(fund)
|
|
),
|
|
chat_id=call.message.chat.id,
|
|
message_id=call.message.id,
|
|
parse_mode="html",
|
|
)
|
|
elif call.data == "no":
|
|
await bot.edit_message_text(
|
|
text=get_fund_text(fund),
|
|
chat_id=call.message.chat.id,
|
|
message_id=call.message.id,
|
|
reply_markup=keyboards.fund_markup(),
|
|
parse_mode="html",
|
|
)
|
|
else:
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.not_owner.format(owner_name=fund.owner.name),
|
|
)
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=call.message.chat.id,
|
|
text=textbook.fund_not_found,
|
|
)
|
|
|
|
|
|
@bot.callback_query_handler(func=lambda c: c.data == "contributed")
|
|
async def contributed(call: types.CallbackQuery):
|
|
group = get_group(call.message.chat)
|
|
user = get_user(call.from_user)
|
|
if fund := group.funds.filter(Fund.active).first():
|
|
if fund_user := fund.members.filter(
|
|
FundMember.user_id == call.from_user.id
|
|
).first():
|
|
if not fund_user.contributed:
|
|
setattr(fund_user, "contributed", True)
|
|
session.commit()
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.contributed,
|
|
)
|
|
await bot.edit_message_text(
|
|
text=get_fund_text(fund),
|
|
chat_id=call.message.chat.id,
|
|
message_id=call.message.id,
|
|
reply_markup=keyboards.fund_markup(),
|
|
parse_mode="html",
|
|
)
|
|
if (
|
|
fund.members.count()
|
|
== fund.members.filter(FundMember.contributed == True).count()
|
|
):
|
|
setattr(fund, "active", False)
|
|
session.commit()
|
|
await bot.edit_message_text(
|
|
text=get_fund_text(fund),
|
|
chat_id=call.message.chat.id,
|
|
message_id=call.message.id,
|
|
reply_markup=None,
|
|
parse_mode="html",
|
|
)
|
|
await bot.send_message(
|
|
chat_id=call.message.chat.id,
|
|
text=textbook.fund_completed.format(
|
|
fund_name=fund.name,
|
|
owner_str=f'<a href="tg://user?id={fund.owner.id}">{fund.owner.name}</a>',
|
|
),
|
|
parse_mode="html",
|
|
)
|
|
else:
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.already_contributed,
|
|
show_alert=True,
|
|
)
|
|
|
|
else:
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id,
|
|
text=textbook.not_fund_member,
|
|
show_alert=True,
|
|
)
|
|
else:
|
|
await bot.answer_callback_query(
|
|
callback_query_id=call.id, text=textbook.fund_not_found, show_alert=True
|
|
)
|
|
|
|
|
|
@bot.message_handler(commands=["remind"])
|
|
async def remind(msg: Message):
|
|
if msg.chat.type in ("group", "supergroup"):
|
|
group = get_group(msg.chat)
|
|
if (
|
|
fund := session.query(Fund)
|
|
.filter(Fund.group_id == group.id, Fund.active == True)
|
|
.first()
|
|
):
|
|
not_contributed = fund.members.filter(FundMember.contributed == False).all()
|
|
if len(not_contributed):
|
|
s = ""
|
|
for fm in not_contributed:
|
|
s += f'<a href="tg://user?id={fm.user.id}">{fm.user.name}</a>\n'
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.remind.format(fund_name=fund.name, s=s),
|
|
parse_mode="html",
|
|
)
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.remind_already.format(fund_name=fund.name),
|
|
parse_mode="html",
|
|
)
|
|
else:
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.fund_not_found,
|
|
)
|
|
|
|
|
|
@bot.message_handler(commands=["mystate"])
|
|
async def mystate(msg: Message):
|
|
state = await bot.get_state(user_id=msg.from_user.id)
|
|
await bot.send_message(chat_id=msg.chat.id, text=state)
|
|
|
|
|
|
@bot.message_handler(commands=["chatid"])
|
|
async def chatid(msg: Message):
|
|
await bot.send_message(chat_id=msg.chat.id, text=msg.chat.id)
|
|
|
|
|
|
async def main():
|
|
a = asyncio.create_task(bot.polling(non_stop=True))
|
|
await a
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("Bot started", flush=True)
|
|
Base.metadata.create_all(engine)
|
|
session = Session()
|
|
bot.add_custom_filter(StateFilter(bot))
|
|
bot.enable_saving_states(filename="./.state-save/states.pkl")
|
|
asyncio.run(main())
|