585 lines
19 KiB
Python
585 lines
19 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
|
|
from sqlalchemy.orm import joinedload
|
|
|
|
|
|
bot = AsyncTeleBot(token, state_storage=StatePickleStorage())
|
|
session: Session = None
|
|
|
|
|
|
class States(StatesGroup):
|
|
default = State()
|
|
newfund_amount = State()
|
|
newfund_description = 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,
|
|
description=fund.description,
|
|
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)
|
|
fund_members_count = (
|
|
session.query(FundMember)
|
|
.join(Fund)
|
|
.filter(Fund.owner_id == user.id, Fund.active == True)
|
|
.options(joinedload(FundMember.fund))
|
|
.distinct(FundMember.fund)
|
|
.count()
|
|
)
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.start_private.format(count=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_list"])
|
|
# async def setup_list(msg: Message):
|
|
# if msg.chat.type in ("group", "supergroup"):
|
|
# group = get_group(msg.chat)
|
|
# members = group.group_members
|
|
# await bot.send_message(
|
|
# chat_id=msg.chat.id,
|
|
# text=textbook.setup_list.format(
|
|
# count=members.count(),
|
|
# members=", ".join(m.user.name for m in members.all()),
|
|
# ),
|
|
# parse_mode="html",
|
|
# )
|
|
|
|
|
|
def get_setup_text(group):
|
|
members = group.group_members
|
|
return textbook.setup + textbook.setup_list.format(count=members.count(), members=", ".join('<a href="tg://user?id={m.user.id}">{m.user.name}</a>\n' for m in members.all()))
|
|
|
|
|
|
@bot.message_handler(commands=["setup"])
|
|
async def setup(msg: Message):
|
|
group = get_group(msg.chat)
|
|
if msg.chat.type in ("group", "supergroup"):
|
|
members = group.group_members
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=get_setup_text(group),
|
|
reply_markup=keyboards.setup(),
|
|
parse_mode="html",
|
|
)
|
|
|
|
|
|
@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)
|
|
group = get_group(msg.chat)
|
|
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.edit_message_text(
|
|
chat_id=msg.chat.id,
|
|
text=get_setup_text(group),
|
|
reply_markup=keyboards.setup(),
|
|
message_id=msg.id,
|
|
parse_mode="html",
|
|
)
|
|
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.callback_query_handler(
|
|
func=lambda c: c.data == "cancel", state=States.newfund_description
|
|
)
|
|
async def cancel_newfund_description(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 msg.text.isdigit():
|
|
async with bot.retrieve_data(
|
|
user_id=msg.from_user.id, chat_id=msg.chat.id
|
|
) as state_data:
|
|
state_data["newfund_amount"] = int(msg.text)
|
|
await bot.set_state(
|
|
user_id=msg.from_user.id,
|
|
chat_id=msg.chat.id,
|
|
state=States.newfund_description,
|
|
)
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=textbook.newfund_description,
|
|
reply_markup=keyboards.cancel(),
|
|
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(content_types=["text"], state=States.newfund_description)
|
|
async def newfund_description(msg: Message):
|
|
if len(msg.text) < 121:
|
|
await bot.set_state(
|
|
user_id=msg.from_user.id, chat_id=msg.chat.id, state=States.default
|
|
)
|
|
async with bot.retrieve_data(
|
|
user_id=msg.from_user.id, chat_id=msg.chat.id
|
|
) as state_data:
|
|
amount = state_data.get("newfund_amount")
|
|
fund = Fund(
|
|
owner_id=msg.from_user.id,
|
|
group_id=msg.chat.id,
|
|
description=msg.text,
|
|
amount=amount,
|
|
name=datetime.today().strftime("%d.%m"),
|
|
)
|
|
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.newfund_description_too_long,
|
|
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=["dmremind"])
|
|
async def dmremind(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):
|
|
counter = 0
|
|
not_sent = []
|
|
for member in not_contributed:
|
|
try:
|
|
await asyncio.sleep(0.1)
|
|
await bot.send_message(
|
|
chat_id=member.user.id,
|
|
text=textbook.dmremind.format(
|
|
fund_name=fund.name, chat_name=msg.chat.title
|
|
),
|
|
)
|
|
counter += 1
|
|
except Exception:
|
|
not_sent.append(member)
|
|
s = textbook.dmremind_completed.format(
|
|
sent_count=counter, members_count=len(not_contributed)
|
|
)
|
|
|
|
if not_sent:
|
|
s += "\n" + textbook.dmremind_not_sent_list.format(
|
|
members=", ".join([m.user.name for m in not_sent])
|
|
)
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id,
|
|
text=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)
|
|
|
|
|
|
@bot.message_handler(commands=["guide"])
|
|
async def guide(msg: Message):
|
|
await bot.send_message(chat_id=msg.chat.id, text=textbook.guide, parse_mode="html")
|
|
|
|
|
|
@bot.message_handler(commands=["commands"])
|
|
async def commands(msg: Message):
|
|
await bot.send_message(
|
|
chat_id=msg.chat.id, text=textbook.commands, parse_mode="html"
|
|
)
|
|
|
|
|
|
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())
|