-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
358 lines (314 loc) · 20 KB
/
main.py
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
import discord
import discord.ext
import random
import math
from datetime import datetime, UTC
from dotenv import load_dotenv, dotenv_values
import os, json, re
from time import sleep
from tinydb import TinyDB, Query, operations
from io import StringIO
from conf.tickettranscripts import *
from html import escape
ignored_keys = {'role_id', 'category_id', 'ticket_count'}
with open('challenges.json','r') as f:
CHALLENGES=json.load(f)
db = TinyDB('db.json')
query = Query()
def remove_ignored_keys(data, keys_to_ignore):
if isinstance(data, list):
return [i for i in [remove_ignored_keys(item, keys_to_ignore) for item in data] if i]
elif isinstance(data, dict):
return {k: remove_ignored_keys(v, keys_to_ignore) for k, v in data.items() if k not in keys_to_ignore}
else:
return data
challenges_data_cleaned = remove_ignored_keys(CHALLENGES, ignored_keys)
db_data_cleaned = remove_ignored_keys(db.all(), ignored_keys)
if not db.all():
db.insert_multiple(CHALLENGES)
db_data_cleaned = remove_ignored_keys(db.all(), ignored_keys)
if not challenges_data_cleaned == db_data_cleaned:
print('---\n\nWARNING: If you have changed challenges.json and you have ran this bot before, delete db.json, there will be errors otherwise. This will also erase your current structure, ids and ticket count so make sure to run /rmsetup beforehand.\n\n---')
def sanitize_html(html_string):
return escape(html_string, quote=True)
def convert_bytes_to_best_size(bytes_num):
suffixes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB']
index = 0
while bytes_num >= 1024 and index < len(suffixes) - 1:
bytes_num /= 1024.0
index += 1
return "{:.1f} {}".format(bytes_num, suffixes[index])
def html_beautify(text):
if '```' in text:
code_blocks = text.split('```')
formatted_text = ""
for i in range(1, len(code_blocks), 2):
lang = code_blocks[i].split('\n')[0].strip()
code_content = code_blocks[i].split('\n')
code_content[0] = ""
code_content = '\n'.join(code_content).strip()
if lang == "":
formatted_code_block = f'<pre class="codeblock" is="message-codeblock"><div class="shitcode"><div class="lines"></div><code class="hljs">{code_content}</code></div><div class="copy">Copy</div></pre>'
else:
formatted_code_block = f'<pre class="codeblock" is="message-codeblock"><div class="lang">{lang}</div><div class="shitcode"><div class="lines"></div><code class="hljs {lang}">{code_content}</code></div><div class="copy">Copy</div></pre>'
formatted_text += formatted_code_block
text = formatted_text
if '~~' in text:
pattern = r'~~(.*?)~~'
text = re.sub(pattern, r'<s>\1</s>', text, flags=re.DOTALL)
if '**' in text:
pattern = r'\*\*(.*?)\*\*'
text = re.sub(pattern, r'<b>\1</b>', text, flags=re.DOTALL)
if '__' in text:
pattern = r'\_\_(.*?)\_\_'
text = re.sub(pattern, r'<ins>\1</ins>', text, flags=re.DOTALL)
if '*' in text:
pattern = r'\*(.*?)\*'
text = re.sub(pattern, r'<i>\1</i>', text, flags=re.DOTALL)
if '_' in text:
pattern = r'\_(.*?)\_'
text = re.sub(pattern, r'<i>\1</i>', text, flags=re.DOTALL)
if '`' in text:
pattern = r'`(.*?)`'
text = re.sub(pattern, r'<code>\1</code>', text, flags=re.DOTALL)
return text
load_dotenv()
TOKEN = os.getenv("TOKEN")
GUILD_ID = int(os.getenv("GUILD_ID"))
ORGANISER_ROLE_ID = int(os.getenv("ORGANISER_ROLE_ID"))
DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY = bool(int(os.getenv("DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY")))
DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY_FORMAT = os.getenv("DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY_FORMAT")
BOT_ROLE_ID = 0
description = '''A CTF-specific Ticketing Bot.
Made by se.al / sealldeveloper'''
def clean_string(input_string):
cleaned_string = re.sub(r'[^a-z0-9- ]+', '', input_string)
return cleaned_string
intents = discord.Intents.default()
client = discord.Client(intents=intents)
tree = discord.app_commands.CommandTree(client)
@tree.command(name="setup", description="Using the configuration provided, create all required roles/categories and setup DB.",guild=discord.Object(id=GUILD_ID))
async def slash_command(interaction: discord.Interaction):
await interaction.response.defer(ephemeral=True)
print(f'[{interaction.user.id}] - ran {interaction.command.name}')
guild = interaction.guild
organiserrole = discord.utils.get(guild.roles, id=ORGANISER_ROLE_ID)
if interaction.user.guild_permissions.administrator or organiserrole in interaction.user.roles:
botrole = discord.utils.get(guild.roles, id=BOT_ROLE_ID)
for cat in range(len(CHALLENGES)):
await interaction.edit_original_response(content=f'Working on {CHALLENGES[cat]["name"]}...')
database_element = db.all()[cat]
if 'role_id' in database_element or 'category_id' in database_element:
return await interaction.edit_original_response(content=f'Run `/rmsetup` first before re-running setup!')
try:
role = await guild.create_role(name=CHALLENGES[cat]["name"])
col = discord.Color(5860729)
await role.edit(color=col)
db.update({'role_id':role.id},doc_ids=[database_element.doc_id])
except:
return await interaction.edit_original_response(content=f'Could not create/manage new role for challenge category `{CHALLENGES[cat]["name"]}`!')
if DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY and DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY_FORMAT:
catname = DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY_FORMAT.replace('<CATEGORY>',CHALLENGES[cat]["name"],1)
try:
category = await guild.create_category(catname)
db.update({'category_id':category.id},doc_ids=[database_element.doc_id])
except:
return await interaction.edit_original_response(content=f'Could not create the new category for challenge category `{catname}`!')
try:
await category.set_permissions(botrole, read_messages=True, send_messages=True)
await category.set_permissions(guild.default_role, read_messages=False)
await category.set_permissions(organiserrole, read_messages=True, send_messages=True)
except:
return await interaction.edit_original_response(content=f'Could not configure permissions for the new category for challenge category `{catname}`!')
if not DISCORD_CATEGORY_PER_CHALLENGE_CATEGORY:
await interaction.edit_original_response(content=f'Creating Tickets category...')
try:
for d in db.all():
if 'category_id' in d:
return await interaction.edit_original_response(content=f'Run `/rmsetup` first before re-running setup!')
category = await guild.create_category('Tickets')
for d in db.all():
db.update({'category_id':category.id},doc_ids=[d.doc_id])
except:
return await interaction.edit_original_response(content=f'Could not create the new category `Tickets`!')
view = Ticket()
await interaction.edit_original_response(content=f'Setup complete!')
await interaction.followup.send(content="Press the button to open the modal with dropdown menus!", view=view)
else:
return await interaction.edit_original_response(content='You are not an admin!')
@tree.command(name="rmsetup", description="Remove the current setup.",guild=discord.Object(id=GUILD_ID))
async def slash_command(interaction: discord.Interaction):
await interaction.response.defer(ephemeral=True)
print(f'[{interaction.user.id}] - ran {interaction.command.name}')
guild = interaction.guild
organiserrole = discord.utils.get(guild.roles, id=ORGANISER_ROLE_ID)
if interaction.user.guild_permissions.administrator or organiserrole in interaction.user.roles:
for d in db.all():
if "name" in d.keys():
await interaction.edit_original_response(content=f'Working on {d["name"]}...')
if 'category_id' in d.keys():
category = discord.utils.get(guild.categories, id=d['category_id'])
if category:
for channel in category.channels:
try:
await channel.delete()
except:
return await interaction.edit_original_response(content=f'Ticket Category `{d["category_id"]}` could not delete the channel <#{channel.id}>.')
try:
await category.delete()
except:
return await interaction.edit_original_response(content=f'Ticket Category `{d["category_id"]}` could not be deleted.')
db.update(operations.delete('category_id'),doc_ids=[d.doc_id])
if 'role_id' in d.keys():
role = discord.utils.get(guild.roles, id=d['role_id'])
if role:
try:
await role.delete()
except:
return await interaction.edit_original_response(content=f'Ticket Role `{d["role_id"]}` could not be deleted.')
db.update(operations.delete('role_id'),doc_ids=[d.doc_id])
db.update({'ticket_count': 0})
await interaction.edit_original_response(content=f'Setup removal complete!')
else:
return await interaction.edit_original_response(content='You are not an admin!')
@client.event
async def on_ready():
print(f'Syncing trees...')
await tree.sync(guild=discord.Object(id=GUILD_ID))
print(f'Finding the bots role...')
guild = discord.utils.get(client.guilds, id=GUILD_ID)
for role in guild.roles:
if len(role.members) == 1 and role.members[0].id == client.user.id and role.permissions.administrator:
global BOT_ROLE_ID
BOT_ROLE_ID = role.id
break
if BOT_ROLE_ID == 0:
print('---\n\nWARNING: No bot role found! Please make sure there is a role unique to the bot with Admin permissions! Things will break!\n\n---')
print('Registering modal views...')
client.add_view(Ticket())
client.add_view(CategorySelectView())
client.add_view(CloseTicketView())
client.add_view(SubOptionSelectView([discord.SelectOption(label="placeholder", description="placeholder")],""))
print(f'Logged in as {client.user} (ID: {client.user.id})')
print('------')
class CloseTicketView(discord.ui.View):
def __init__(self):
super().__init__(timeout=None)
@discord.ui.button(label="Close Ticket", style=discord.ButtonStyle.danger, custom_id="close_ticket")
async def close_ticket(self, interaction: discord.Interaction, button: discord.ui.Button):
transcript = StringIO()
transcript.write(f"""<!DOCTYPE html><html lang="en"><head><meta charset="UTF-8"><meta http-equiv="X-UA-Compatible" content="ie=edge"><meta name="viewport" content="width=device-width,initial-scale=1,maximum-scale=1,user-scalable=no"><title>#{interaction.channel.name} Transcript</title>""")
transcript.write(TICKET_HEADER)
transcript.write(f"""<span class="name">{interaction.channel.name}</span><theme-switch></theme-switch></header><main><div class="placeholder"></div><div class="welcome"><h1>Welcome to <b>#{interaction.channel.name}</b>!</h1><div>This is the start of the #{interaction.channel.name} channel.</div></div><discord-messages>""")
messages = []
async for message in interaction.channel.history(limit=None):
timestamp = datetime.fromtimestamp(message.created_at.timestamp(), UTC)
message.content = html_beautify(sanitize_html(message.content)).replace('\n','<br>')
to_append = ""
to_append += f"""<discord-message data-id="0" data-author="{message.author.id}" class="group-start"><img is="message-avatar" src="{message.author.display_avatar.with_size(512).url}" data-discriminator="0000" alt="avatar" class="avatar" onerror="this.onerror=null;this.src='https://cdn.discordapp.com/embed/avatars/0.png';"><message-date class="time" data-type="time" data-timestamp="{message.created_at}">{timestamp.strftime('%H:%M')}</message-date><div class="contents"><message-header><span class="name">{message.author.display_name}</span><span class="badge"></span><message-date class="date" data-type="date" data-timestamp="{message.created_at}">{timestamp.strftime('%a, %d %b %Y %H:%M:%S GMT')}</message-date></message-header><message-markup>{message.content}</message-markup>"""
if len(message.attachments) > 0:
for attachment in message.attachments:
to_append += f"""<message-attachment class><div class="data"><img src="https://discord.com/assets/985ea67d2edab4424c62009886f12e44.svg" alt class="icon"><div class="details"><a href="{attachment.url}" target="_blank">{attachment.filename}</a><span>{convert_bytes_to_best_size(attachment.size)}</span></div><a href="{attachment.url}" target="_blank" class="download"><svg width="24" height="24" viewBox="0 0 24 24"><path fill="currentColor" d="M19 9h-4V3H9v6H5l7 7 7-7zM5 18v2h14v-2H5z" /></svg></a></div></message-attachment>"""
to_append += """</div>"""
to_append += """</div></discord-message>"""
messages.append(to_append)
for x in reversed(messages):
transcript.write(x)
createdtimestamp = datetime.fromtimestamp(message.created_at.timestamp(), UTC)
transcript.write(f"""</discord-messages></main><footer><span>This archive was generated on the <message-date data-type="full" data-timestamp="{interaction.created_at}">{createdtimestamp.strftime('%d %b %Y at %H:%M:%S Coordinated Universal Time')}</message-date></span></footer></body></html>""")
transcript.seek(0)
file = discord.File(transcript, filename="ticket_transcript.html")
overwrites = interaction.channel.overwrites
user = None
for target, overwrite in overwrites.items():
if isinstance(target,discord.object.Object):
user = await interaction.guild.fetch_member(target.id)
await user.send("Here is the transcript of your ticket:", file=file)
await interaction.response.send_message("Transcript of the ticket sent to your DMs! Deleting channel in 10 seconds...", ephemeral=True)
sleep(10)
return await interaction.channel.delete()
class Ticket(discord.ui.View):
def __init__(self):
super().__init__(timeout=None)
@discord.ui.button(label="Create a ticket", style=discord.ButtonStyle.primary, custom_id="ticket_create",emoji="🎫")
async def button_callback(self, interaction: discord.Interaction, button: discord.ui.Button):
await interaction.response.send_message("Select a challenge category:", view=CategorySelectView(), ephemeral=True)
class CategorySelectView(discord.ui.View):
def __init__(self, *args, **kwargs) -> None:
super().__init__(timeout=None)
self.add_item(CategorySelect())
async def on_timeout(self):
pass
class CategorySelect(discord.ui.Select):
def __init__(self, *args, **kwargs):
options = []
for cat in CHALLENGES:
options.append(discord.SelectOption(label=cat["name"], description=f"The challenges in the {cat['name']} category."))
super().__init__(placeholder="Challenge Category", options=options, custom_id="category_select", *args, **kwargs)
async def callback(self, interaction: discord.Interaction):
selected_option = self.values[0]
challenge_options = []
for cat in CHALLENGES:
if cat['name'] == selected_option:
for chal in cat['challenges']:
challenge_options.append(
discord.SelectOption(label=chal['name'], description=f"The challenge {chal['name']} in the {cat['name']} category.")
)
break
await interaction.response.edit_message(content=f"Selected Category: `{selected_option}`.", view=SubOptionSelectView(challenge_options, selected_option))
class SubOptionSelectView(discord.ui.View):
def __init__(self, options, selected_category, *args, **kwargs) -> None:
super().__init__(timeout=None)
self.add_item(SubOptionSelect(options, selected_category))
self.add_item(GoBackButton(selected_category))
async def on_timeout(self):
pass
class GoBackButton(discord.ui.Button):
def __init__(self, selected_category):
super().__init__(label="Go Back", style=discord.ButtonStyle.secondary, custom_id="go_back", row=1)
self.selected_category = selected_category
async def callback(self, interaction: discord.Interaction):
await interaction.response.edit_message(content=f"Select a challenge category:", view=CategorySelectView())
class SubOptionSelect(discord.ui.Select):
def __init__(self, options, selected_category, *args, **kwargs):
super().__init__(placeholder="Select a challenge:", options=options, custom_id="challenge_select", *args, **kwargs)
self.selected_category = selected_category
async def callback(self, interaction: discord.Interaction):
selected_option = self.values[0]
selected_category = self.selected_category
await interaction.response.edit_message(content=f"You selected: `{selected_option}` from the category `{selected_category}`. Creating ticket...", view=None)
guild = discord.utils.get(client.guilds, id=GUILD_ID)
for cat in range(len(CHALLENGES)):
if selected_category == CHALLENGES[cat]['name']:
db_elem = db.all()[cat]
category = discord.utils.get(guild.categories, id=db_elem['category_id'])
if not db.search(query.ticket_count.exists()):
db.insert({'ticket_count': 0})
db.update({'ticket_count': db.get(query.ticket_count.exists())['ticket_count'] + 1})
ticket_number = db.get(query.ticket_count.exists())['ticket_count']
ticket = await category.create_text_channel(name=f'ticket-{ticket_number}')
await ticket.set_permissions(interaction.user, read_messages=True, send_messages=True)
initial_message = f"Ticket about challenge `{selected_option}` from the category `{selected_category}`. The ticket belongs to <@{interaction.user.id}> (`{interaction.user.id}`)."
if 'ping_creators' in db_elem.keys() or 'ping_category' in db_elem.keys():
initial_message+="\n\nPing(s): "
if 'ping_creators' in db_elem.keys():
if db_elem['ping_creators'] == True:
for chal in db_elem['challenges']:
if chal['name'] == selected_option:
if 'creators' in chal.keys():
if len(chal['creators']) == 0:
initial_message += "Nobody to ping!"
else:
for creator in chal['creators']:
initial_message += f"<@{creator}> "
if 'ping_category' in db_elem.keys():
if db_elem['ping_category'] == True:
role_id = db_elem['role_id']
initial_message += f"<@&{role_id}>"
close_ticket = CloseTicketView()
await ticket.send(initial_message,view=close_ticket)
await interaction.edit_original_response(content=f"Your ticket has been made! See it here: <#{ticket.id}>")
sleep(5)
return await interaction.delete_original_response()
# run the bot
client.run(TOKEN)