new functions added

added next functions: -delete event -set current event -current flags -delete invalid flag -show current event
Beadyhx · May 21, 2023 · 8fb51aa · 8fb51aa
1 parent c4b4e24
commit 8fb51aa
Showing 1 changed file with 243 additions and 23 deletions.
diff --git a/flagbot.py b/flagbot.py
@@ -29,6 +29,7 @@ class MyStates(StatesGroup):
     task = State()
     flag = State() # statesgroup should contain states
     event = State()
+    template = State()
 
 def log(text, filename='log.txt'):
     with open(filename, 'a') as logfile:
@@ -39,11 +40,38 @@ def cancelbutton():
     keyboard.add(telebot.types.InlineKeyboardButton(text="Отмена", callback_data="/cancel"))
     return keyboard
 
-@bot.message_handler(state=MyStates.task)
-async def task_get(message):
+@bot.message_handler(state=MyStates.event)
+async def event_get(message):
     """
     State 1. Will process when user's state is MyStates.task.
     """
+    keyboard = cancelbutton()
+    await bot.send_message(message.chat.id, f'Введите образец флага, наример SomeCTF{{}}:', reply_markup=keyboard)
+    await bot.set_state(message.from_user.id, MyStates.template, message.chat.id)
+    async with bot.retrieve_data(message.from_user.id, message.chat.id) as data:
+        data['event'] = message.text
+
+@bot.message_handler(state=MyStates.template)
+async def template_get(message):
+    template = message.text
+    async with bot.retrieve_data(message.from_user.id, message.chat.id) as data:
+        event = data['event']
+    setzero_query = "UPDATE events SET iscurrent=0"
+    insert_query = "INSERT INTO events(event, template, iscurrent) VALUES (%s, %s, %s)"
+    insert_values = (event, template, 1)
+    cursor.execute(setzero_query)
+    cursor.execute(insert_query, insert_values)
+    db.commit()
+    # Send a confirmation to the user
+    username = message.from_user.username
+    async with bot.retrieve_data(message.from_user.id, message.chat.id) as data:
+        await bot.send_message(message.chat.id, "Новый ивент создан и назначен текущим.")
+    await bot.delete_state(message.from_user.id, message.chat.id)
+    log('template_get:'+' '+str(username)+' '+str(event)+' '+str(template))
+    await start(message)
+
+@bot.message_handler(state=MyStates.task)
+async def task_get(message):
     keyboard = cancelbutton()
     await bot.send_message(message.chat.id, f'Введите флаг:', reply_markup=keyboard)
     await bot.set_state(message.from_user.id, MyStates.flag, message.chat.id)
@@ -58,9 +86,14 @@ async def flag_get(message):
     flag = message.text
     async with bot.retrieve_data(message.from_user.id, message.chat.id) as data:
         task = data['task']
-    cursor.execute("SELECT event FROM events WHERE iscurrent = 1")
-    current_event = cursor.fetchone()[0]
-    # Get user data from the callback query
+    try:
+        cursor.execute("SELECT event FROM events WHERE iscurrent = 1")
+        current_event = cursor.fetchone()[0]
+    except:
+        await bot.send_message(message.chat.id, "Ошибка: Отсутствует текущий ивент. Это означает, что в данный момент в команде не проходит CTF.\nЕсли это не так, обратитесь к администратору c просьбой добавить новый ивент или сделать текущим один из существующих.")
+        await bot.delete_state(message.from_user.id, message.chat.id)
+        await start(message)
+        return 0
     username = message.from_user.username
     submit_time = gettime()
     insert_query = "INSERT INTO flags(username, event, task, flag, submit_time) VALUES (%s, %s, %s, %s, %s)"
@@ -83,13 +116,20 @@ def startkeyboard(username):
     keyboard = telebot.types.InlineKeyboardMarkup()
     if username in admin_list:
         keyboard.add(telebot.types.InlineKeyboardButton(text="Добавить флаг", callback_data="/addflag"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Текущий ивент", callback_data="/showcurrent"))
         keyboard.add(telebot.types.InlineKeyboardButton(text="Ивенты", callback_data="/events"))
-        keyboard.add(telebot.types.InlineKeyboardButton(text="Добавить ивент (пока что не работает)", callback_data="/addevent"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Добавить текущий ивент", callback_data="/addevent"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Сделать ивент текущим", callback_data="/setcurrent"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Удалить ивент", callback_data="/deleteevent"))
         keyboard.add(telebot.types.InlineKeyboardButton(text="Мои флаги", callback_data="/myflags"))
         keyboard.add(telebot.types.InlineKeyboardButton(text="Все флаги", callback_data="/flags"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Текущие флаги", callback_data="/currentflags"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Удалить невалидный флаг", callback_data="/deleteflag"))
     else:
         keyboard.add(telebot.types.InlineKeyboardButton(text="Добавить флаг", callback_data="/addflag"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Текущий ивент", callback_data="/showcurrent"))
         keyboard.add(telebot.types.InlineKeyboardButton(text="Мои флаги", callback_data="/myflags"))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Удалить невалидный флаг", callback_data="/deleteflag"))
     return keyboard
 
 async def start(message):
@@ -131,8 +171,11 @@ async def showevents(call):
         await bot.send_message(call.message.chat.id, 'Недостаточно прав')
     await startcallback(call)
 
-def addevent(call):
-    pass
+async def addevent(call):
+    await bot.set_state(call.message.chat.id, MyStates.event, call.message.chat.id)
+    keyboard = cancelbutton()
+    await bot.send_message(call.message.chat.id, 'ВНИМАНИЕ: Новый ивент сразу будет сделан текущим. Если вы не уверены, нажмите \'Отмена\'.\nВведите название ивента:', reply_markup=keyboard)
+    # Select the current event from the events table
 
 @bot.message_handler(state=MyStates.event)
 async def event_get(message):
@@ -142,7 +185,29 @@ async def event_get(message):
     async with bot.retrieve_data(message.from_user.id, message.chat.id) as data:
         await bot.send_message(message.chat.id, message.text)
 
-
+async def showcurrentflags(call):
+    username = call.from_user.username
+    if username in admin_list:
+        cursor = db.cursor()
+        cursor.execute("SELECT event FROM events WHERE iscurrent = 1")
+        current_event = cursor.fetchone()[0]
+        select_query = "SELECT event, flag, task, username, submit_time FROM flags WHERE event = %s"
+        select_values = (current_event,)
+        cursor.execute(select_query, select_values)
+        flags = cursor.fetchall()
+        cursor.close()
+        response_message = ''
+        for flag in flags:
+            string = ''
+            for i in range(0,len(flag)):
+                string += str(flag[i])+' # '
+            response_message += string+'\n'
+            response_message += '################\n'
+        await bot.send_message(call.message.chat.id, response_message)
+        log('showcurrentflags:'+' '+str(username)+' '+str(flags))
+    else:
+        await bot.send_message(call.message.chat.id, "Недостаточно прав.")
+    await startcallback(call)
 
 async def showmyflags(call):
     username = call.from_user.username
@@ -167,26 +232,95 @@ async def showallflags(call):
     username = call.from_user.username
     if username in admin_list:
         cursor = db.cursor()
-        cursor.execute("SELECT event FROM events WHERE iscurrent = 1")
-        current_event = cursor.fetchone()[0]
-        select_query = "SELECT event, flag, task, username, submit_time FROM flags WHERE event = %s"
-        select_values = (current_event,)
+        select_query = "SELECT DISTINCT event FROM flags"
+        cursor.execute(select_query)
+        events = cursor.fetchall()
+        cursor.close()
+        keyboard = telebot.types.InlineKeyboardMarkup()
+        log('showallflags:'+' '+str(username)+' '+str(events))
+        # Create a list of buttons for each event
+        for event in events:
+            keyboard.add(telebot.types.InlineKeyboardButton(text=str(event[0]), callback_data=str('showallflags'+'$##$'+event[0])))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Отмена", callback_data="/cancel"))
+        # Create a custom keyboard with the event buttons and send
+        await bot.send_message(call.message.chat.id, 'Выберите ивент:', reply_markup=keyboard)
+    else:
+        await bot.send_message(call.message.chat.id, "Недостаточно прав.")
+        await startcallback(call)
+
+async def deleteevent(call):
+    username = call.from_user.username
+    if username in admin_list:
+        cursor = db.cursor()
+        select_query = "SELECT DISTINCT event FROM events"
+        cursor.execute(select_query)
+        events = cursor.fetchall()
+        cursor.close()
+        keyboard = telebot.types.InlineKeyboardMarkup()
+        log('deleteevent:'+' '+str(username)+' '+str(events))
+        # Create a list of buttons for each event
+        for event in events:
+            keyboard.add(telebot.types.InlineKeyboardButton(text=str(event[0]), callback_data=str('deleteevent'+'$##$'+event[0])))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Отмена", callback_data="/cancel"))
+        # Create a custom keyboard with the event buttons and send
+        await bot.send_message(call.message.chat.id, 'Выберите ивент:', reply_markup=keyboard)
+    else:
+        await bot.send_message(call.message.chat.id, "Недостаточно прав.")
+        await startcallback(call)
+
+async def deleteflag(call):
+    username = call.from_user.username
+    if username in admin_list:
+        cursor = db.cursor()
+        select_query = "SELECT flag FROM flags"
+        cursor.execute(select_query)
+        flags = cursor.fetchall()
+        print(flags)
+        cursor.close()
+        keyboard = telebot.types.InlineKeyboardMarkup()
+        log('deleteflag:'+' '+str(username)+' '+str(flags))
+        # Create a list of buttons for each event
+        for flag in flags:
+            keyboard.add(telebot.types.InlineKeyboardButton(text=str(flag[0]), callback_data=str('deleteflag'+'$##$'+flag[0])))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Отмена", callback_data="/cancel"))
+        # Create a custom keyboard with the event buttons and send
+        await bot.send_message(call.message.chat.id, 'Выберите флаг:', reply_markup=keyboard)
+    else:
+        cursor = db.cursor()
+        select_query = "SELECT flag FROM flags WHERE username = %s"
+        select_values = (username,)
         cursor.execute(select_query, select_values)
         flags = cursor.fetchall()
+        print(flags)
         cursor.close()
-        response_message = ''
+        keyboard = telebot.types.InlineKeyboardMarkup()
+        log('deleteflag:'+' '+str(username)+' '+str(flags))
+        # Create a list of buttons for each event
         for flag in flags:
-            string = ''
-            for i in range(0,len(flag)):
-                string += str(flag[i])+' # '
-            response_message += string+'\n'
-            response_message += '################\n'
-        await bot.send_message(call.message.chat.id, response_message)
-        log('showallflags:'+' '+str(username)+' '+str(flags))
+            keyboard.add(telebot.types.InlineKeyboardButton(text=str(flag[0]), callback_data=str('deleteflag'+'$##$'+flag[0])))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Отмена", callback_data="/cancel"))
+        # Create a custom keyboard with the event buttons and send
+        await bot.send_message(call.message.chat.id, 'Выберите флаг:', reply_markup=keyboard)
+
+async def setcurrent(call):
+    username = call.from_user.username
+    if username in admin_list:
+        cursor = db.cursor()
+        select_query = "SELECT DISTINCT event FROM events"
+        cursor.execute(select_query)
+        events = cursor.fetchall()
+        cursor.close()
+        keyboard = telebot.types.InlineKeyboardMarkup()
+        log('setcurrent:'+' '+str(username)+' '+str(events))
+        # Create a list of buttons for each event
+        for event in events:
+            keyboard.add(telebot.types.InlineKeyboardButton(text=str(event[0]), callback_data=str('setcurrent'+'$##$'+event[0])))
+        keyboard.add(telebot.types.InlineKeyboardButton(text="Отмена", callback_data="/cancel"))
+        # Create a custom keyboard with the event buttons and send
+        await bot.send_message(call.message.chat.id, 'Выберите ивент:', reply_markup=keyboard)
     else:
         await bot.send_message(call.message.chat.id, "Недостаточно прав.")
-    await startcallback(call)
-
+        await startcallback(call)
 
 @bot.message_handler(commands=['start'])
 async def handle_start_command(message):
@@ -229,20 +363,52 @@ async def handle_callback_query(call):
         # Handle "/start" command separately
         await addevent(call)
         return 0
+    elif call.data == "/setcurrent":
+        # Handle "/start" command separately
+        await setcurrent(call)
+        return 0
     elif call.data == "/myflags":
         # Handle "/start" command separately
         await showmyflags(call)
         return 0
+    elif call.data == "/currentflags":
+        # Handle "/start" command separately
+        await showcurrentflags(call)
+        return 0
     elif call.data == "/flags":
         # Handle "/start" command separately
         await showallflags(call)
         return 0
+    elif call.data == "/deleteevent":
+        # Handle "/start" command separately
+        await deleteevent(call)
+        return 0
+    elif call.data == "/deleteflag":
+        # Handle "/start" command separately
+        await deleteflag(call)
+        return 0
     elif call.data == "/cancel":
         # Handle "/start" command separately
         await bot.send_message(call.message.chat.id, "Отмена...")
         await bot.delete_state(call.message.chat.id, call.message.chat.id)
         await startcallback(call)
         return 0
+    elif call.data == "/showcurrent":
+        cursor = db.cursor()
+        get_event_query = "SELECT event FROM events WHERE iscurrent=1"
+        cursor.execute(get_event_query)
+        event = cursor.fetchone()[0]
+        print(event)
+        cursor.close()
+        if len(event) > 1:
+            await bot.send_message(call.message.chat.id, str('Текущий ивент: '+event))
+            await startcallback(call)
+            return 0
+        else:
+            await bot.send_message(call.message.chat.id, "Ошибка: Отсутствует текущий ивент. Это означает, что в данный момент в команде не проходит CTF.\nЕсли это не так, обратитесь к администратору c просьбой добавить новый ивент или сделать текущим один из существующих.")
+            await startcallback(call)
+            return 0
+
     elif call.data in eventlist:
         cursor = db.cursor()
         select_query = "SELECT username, task, flag FROM flags WHERE event = %s AND username = %s"
@@ -259,6 +425,60 @@ async def handle_callback_query(call):
         await bot.send_message(call.message.chat.id, response_message)
         cursor.close()
         await startcallback(call)
+    elif call.data.split('$##$')[0] == 'showallflags':
+        cursor = db.cursor()
+        select_query = "SELECT flag, task, username, submit_time FROM flags WHERE event = %s"
+        select_values = (call.data.split('$##$')[1],)
+        cursor.execute(select_query, select_values)
+        flags = cursor.fetchall()
+        cursor.close()
+        response_message = ''
+        for flag in flags:
+            string = ''
+            for i in range(0,len(flag)):
+                string += str(flag[i])+' # '
+            response_message += string+'\n'
+            response_message += '################\n'
+        log('showcurrentflags:'+' '+str(username)+' '+str(flags))
+        await bot.send_message(call.message.chat.id, response_message)
+        await startcallback(call)
+    elif call.data.split('$##$')[0] == 'deleteevent':
+        cursor = db.cursor()
+        delete_event_query = "DELETE FROM events WHERE event = %s"
+        delete_event_values = (call.data.split('$##$')[1],)
+        cursor.execute(delete_event_query, delete_event_values)
+        # delete_flags_query = "DELETE FROM flags WHERE event = %s"
+        # delete_flags_values = (call.data.split('$##$')[1],)
+        # cursor.execute(delete_flags_query, delete_flags_values)
+        db.commit()
+        cursor.close()
+        response_message = str(call.data.split('$##$')[1]+': ивент успешно удалён')
+        await bot.send_message(call.message.chat.id, response_message)
+        await startcallback(call)
+    elif call.data.split('$##$')[0] == 'deleteflag':
+        for text in call.data.split('$##$'):
+            print(text)
+        cursor = db.cursor()
+        delete_flags_query = "DELETE FROM flags WHERE flag = %s"
+        delete_flags_values = (call.data.split('$##$')[1],)
+        cursor.execute(delete_flags_query, delete_flags_values)
+        db.commit()
+        cursor.close()
+        response_message = str(call.data.split('$##$')[1]+': флаг успешно удалён')
+        await bot.send_message(call.message.chat.id, response_message)
+        await startcallback(call)
+    elif call.data.split('$##$')[0] == 'setcurrent':
+        cursor = db.cursor()
+        setzero_query = "UPDATE events SET iscurrent=0"
+        set_event_query = "UPDATE events SET iscurrent=1 WHERE event = %s"
+        set_event_values = (call.data.split('$##$')[1],)
+        cursor.execute(setzero_query)
+        cursor.execute(set_event_query, set_event_values)
+        db.commit()
+        cursor.close()
+        response_message = str(call.data.split('$##$')[1]+': ивент успешно сделан текущим')
+        await bot.send_message(call.message.chat.id, response_message)
+        await startcallback(call)
 
 bot.add_custom_filter(asyncio_filters.StateFilter(bot))