from telegram.ext import Updater, MessageHandler, Filters, CallbackContext, CommandHandler from telegram import Update, InputMediaPhoto import json import sys import io import logging from PIL import Image import random from mwt import MWT submissions = {} names = {} polls = {} poll_options = {} status = {} @MWT(timeout=60*60) def get_admin_ids(bot, chat_id): """Returns a list of admin IDs for a given chat. Results are cached for 1 hour.""" return [admin.user.id for admin in bot.get_chat_administrators(chat_id)] def ensure_admin(func): def wrapper(update, context, *args, **kwargs): chat_id = update.effective_chat.id user_id = update.effective_user.id if user_id in get_admin_ids(context.bot, chat_id): return func(update, context, *args, **kwargs) return wrapper def ensure_chat(func): def wrapper(update, *args, **kwargs): chat_id = update.effective_chat.id if chat_id not in submissions: submissions[chat_id] = {} names[chat_id] = {} status[chat_id] = "closed" polls[chat_id] = {} poll_options[chat_id] = [] print("Added chat: ", chat_id) return func(update, *args, **kwargs) return wrapper def ensure_user(func): def wrapper(update, *args, **kwargs): chat_id = update.effective_chat.id user_id = update.effective_user.id if user_id not in submissions[chat_id]: names[chat_id][user_id] = update.effective_user.first_name print("Added user: ", user_id, " - ", update.effective_user.first_name) return func(update, *args, **kwargs) return wrapper def update_name(func): def wrapper(update, *args, **kwargs): chat_id = update.effective_chat.id user_id = update.effective_user.id if user_id not in submissions[update.effective_chat.id]: names[chat_id][user_id] = update.effective_user.first_name print("Got name for ", user_id, ": ", update.effective_user.first_name) return func(update, *args, **kwargs) return wrapper def ensure_closed(func): def wrapper(update, context, *args, **kwargs): chat_id = update.effective_chat.id if status[chat_id] == "closed": return func(update, context, *args, **kwargs) else: context.bot.send_message(chat_id=chat_id, text="I need to be closed to do that. I'm currently: " + status[chat_id]) return wrapper def ensure_open(func): def wrapper(update, context, *args, **kwargs): chat_id = update.effective_chat.id if status[chat_id] == "open": return func(update, context, *args, **kwargs) else: context.bot.send_message(chat_id=chat_id, text="I need to be open to do that. I'm currently: " + status[chat_id]) return wrapper def ensure_polling(func): def wrapper(update, context, *args, **kwargs): chat_id = update.effective_chat.id if status[chat_id] == "polling": return func(update, context, *args, **kwargs) else: context.bot.send_message(chat_id=chat_id, text="I need to be polling to do that. I'm currently: " + status[chat_id]) return wrapper @ensure_chat @ensure_user @update_name @ensure_open def accept_photo(update: Update, context: CallbackContext): chat_id = update.effective_chat.id user_id = update.effective_user.id file_id = update.message.photo[-1].file_id submissions[chat_id][user_id] = file_id print(file_id) @ensure_chat @ensure_admin @ensure_closed def open_submissions(update: Update, context: CallbackContext): chat_id = update.effective_chat.id submissions[chat_id] = {} context.bot.send_message(chat_id=chat_id, text="Send your images as a reply to this post.") status[chat_id] = "open" @ensure_chat @ensure_admin @ensure_open def start_poll(update: Update, context: CallbackContext): chat_id = update.effective_chat.id status[chat_id] = "polling" poll_options[chat_id] = [] users = list(submissions[chat_id].keys()) print("Users with submissions: ", users) chosen_users = random.sample(users, k=min(len(users), 10)) options = [] chosen = [] for index, user_id in enumerate(chosen_users): options.append(str(index+1) + ": " + names[chat_id][user_id]) chosen.append(submissions[chat_id][user_id]) poll_options[chat_id].append(user_id) if len(chosen) == 0: context.bot.send_message(chat_id=chat_id, text="There were no submissions.") status[chat_id] = "closed" elif len(chosen) == 1: context.bot.send_photo(chat_id=chat_id, photo=chosen[0]) context.bot.send_message(chat_id=chat_id, text="There was only one submission. " + options[0] + " is the winner.") status[chat_id] = "closed" else: input_media = list(map(lambda pair: InputMediaPhoto(pair[1], caption=pair[0]), zip(options, chosen))) context.bot.send_media_group(chat_id=chat_id, media=input_media) polls[chat_id] = context.bot.send_poll(chat_id=chat_id, options = options, question="Pick an icon").message_id @ensure_chat @ensure_admin @ensure_polling def decide_poll(update: Update, context: CallbackContext): chat_id = update.effective_chat.id status[chat_id] = "closed" result = context.bot.stop_poll(chat_id=chat_id, message_id=polls[chat_id]) votes = -1 winners = [] for index, option in enumerate(result.options): if option.voter_count > votes: winners = [] if option.voter_count >= votes: winners.append(index) votes = option.voter_count winner = random.choice(winners) winner_id = poll_options[chat_id][winner] winner_name = names[chat_id][winner_id] winner_photo = submissions[chat_id][winner_id] context.bot.send_message(chat_id=chat_id, text="The winner was " + winner_name) context.bot.send_photo(chat_id=chat_id, photo=winner_photo) @ensure_chat @ensure_admin def cancel(update: Update, context: CallbackContext): chat_id = update.effective_chat.id status[chat_id] = "closed" submissions[chat_id] = {} del polls[chat_id] @ensure_chat def show_status(update: Update, context: CallbackContext): chat_id = update.effective_chat.id if status[chat_id] == "closed": context.bot.send_message(chat_id=chat_id, text="I am currently closed.") elif status[chat_id] == "open": context.bot.send_message(chat_id=chat_id, text=f"I am currently open for submissions. There are {len(submissions[chat_id])} submissions.") elif status[chat_id] == "polling": context.bot.send_message(chat_id=chat_id, text="I am currently polling.") if __name__ == "__main__": try: config = json.load(open("config.json", "r", encoding="utf-8")) except: logging.error("Couldn't read the config!") sys.exit(1) updater = Updater(token=config["token"], use_context=True) dispatcher = updater.dispatcher dispatcher.add_handler(MessageHandler(Filters.photo, accept_photo)) dispatcher.add_handler(CommandHandler("open", open_submissions)) dispatcher.add_handler(CommandHandler("poll", start_poll)) dispatcher.add_handler(CommandHandler("decide", decide_poll)) dispatcher.add_handler(CommandHandler("cancel", cancel)) dispatcher.add_handler(CommandHandler("status", show_status)) updater.start_polling() updater.idle()