guilddump/guilddump.py
2022-05-08 23:26:29 +02:00

298 lines
10 KiB
Python

import asyncio
import csv
import pathlib
import time
import discord
from simple_term_menu import TerminalMenu
async def input_prompt(question: str, default: str = '', checker=None) -> str:
while True:
try:
result = input(question) or default
except KeyboardInterrupt:
exit(1)
if checker is None:
return result
elif asyncio.iscoroutinefunction(checker) and await checker(result):
return result
elif not asyncio.iscoroutinefunction(checker) and checker(result):
return result
print('\r')
async def option_prompt(question: str, options: [str]) -> int:
print(question)
return TerminalMenu(options, clear_menu_on_exit=False).show()
async def select_prompt(question: str, options: [str]) -> [int]:
print(question)
return TerminalMenu(options, clear_menu_on_exit=False, multi_select=True, preselected_entries=options).show()
async def yesno_prompt(question: str) -> bool:
return option_prompt(question, ['Yes', 'No']) == 0
async def request_client() -> discord.Client:
intents = discord.Intents().default()
intents.members = True
async def checker(checker_token: str) -> bool:
checker_client = discord.Client(intents=intents)
try:
await checker_client.login(checker_token)
await checker_client.close()
return True
except discord.LoginFailure:
print('Invalid token')
await checker_client.close()
return False
token = await input_prompt('Discord Bot Token: ', checker=checker)
client = discord.Client(intents=intents)
await client.login(token)
return client
async def request_guild(client: discord.Client) -> discord.Guild:
options = []
guilds = [guild async for guild in client.fetch_guilds(limit=51)]
if len(guilds) > 50:
manual = True
elif len(guilds) <= 10:
manual = False
else:
option = await option_prompt(f'{len(guilds)} guilds are available. Choose from list or enter manually?: ', ['Choose', 'Manually'])
manual = option == 1
if manual:
async def checker(guild_id: str) -> bool:
try:
return await client.fetch_guild(int(guild_id)) is not None
except ValueError:
print('Enter a valid id (must only contain number)')
except discord.Forbidden:
print('Bot not in guild')
return False
return await client.fetch_guild(int(await input_prompt('Guild id to scrape: ', checker=checker)))
else:
for guild in guilds:
options.append(f'{guild.id} ({guild.name})')
return guilds[await option_prompt('Select guild to scrape: ', options)]
async def request_directory() -> pathlib.Path:
async def checker(directory: str) -> bool:
path = pathlib.Path(directory)
if path.exists():
if path.is_dir():
return True
else:
print(f'{directory} is not a directory')
else:
if await yesno_prompt(f'{directory} does not exist. Do you want to create it?'):
try:
path.mkdir(parents=True)
return True
except OSError as e:
print(f'Failed to create directory: {e}')
return False
return pathlib.Path(await input_prompt('Directory to save content in: ', checker=checker))
async def request_channels(guild: discord.Guild) -> [discord.TextChannel]:
options = list(filter(lambda channel: isinstance(channel, discord.TextChannel), await guild.fetch_channels()))
selected = await select_prompt('Select channels to scrape: ', [f'{channel.id} (#{channel.name})' for channel in options])
return [options[select] for select in selected]
async def request_max_members() -> int:
def checker(number: str) -> bool:
try:
if int(number) > 0:
return True
else:
print('Value must be greater than 0')
return False
except ValueError:
return False
return int(await input_prompt(f'Max members to fetch (default: 1000): ', default='1000', checker=checker))
async def request_max_messages() -> int:
def checker(number: str) -> bool:
try:
if int(number) > 0:
return True
else:
print('Value must be greater than 0')
return False
except ValueError:
return False
return int(await input_prompt(f'Max messages to fetch per channel (default: 1000): ', default='1000', checker=checker))
async def dump_channels(guild: discord.Guild, channel_writer):
channel_writer.writerow(['id', 'created_at', 'name', 'type'])
for channel in await guild.fetch_channels():
channel_type = ''
if isinstance(channel, discord.TextChannel):
channel_type = 'text'
elif isinstance(channel, discord.VoiceChannel):
channel_type = 'voice'
elif isinstance(channel, discord.CategoryChannel):
channel_type = 'category'
elif isinstance(channel, discord.StageChannel):
channel_type = 'stage'
channel_writer.writerow([
channel.id,
int(time.mktime(channel.created_at.timetuple())),
channel.name,
channel_type
])
async def dump_members(guild: discord.Guild, limit: int, user_writer):
user_writer.writerow(['id', 'joined_at', 'name', 'nick', 'is_bot', 'premium_since'])
async for member in guild.fetch_members(limit=limit):
user_writer.writerow([
member.id,
int(time.mktime(member.joined_at.timetuple())),
str(member),
member.nick,
member.bot,
int(time.mktime(member.premium_since.timetuple())) if member.premium_since else 0
])
async def dump_roles(guild: discord.Guild, role_writer):
role_writer.writerow(['id', 'created_at', 'name', 'position', 'permissions', 'member_count', 'is_mentionable', 'color'])
for role in await guild.fetch_roles():
role_writer.writerow([
role.id,
role.created_at,
role.name,
role.position,
role.permissions.value,
len(role.members),
role.mentionable,
role.color
])
async def dump_messages(channel: discord.TextChannel, limit: int, message_writer, attachment_writer, reaction_writer, embed_writer):
message_writer.writerow(['id', 'author_id', 'created_at', 'modified_at', 'content'])
attachment_writer.writerow(['message_id', 'type', 'size', 'filename', 'url', 'is_spoiler'])
reaction_writer.writerow(['message_id', 'name', 'reaction_count', 'is_animated'])
embed_writer.writerow(['message_id', 'title', 'description', 'footer', 'image', 'thumbnail', 'video', 'author', 'field_count', 'color'])
async for message in channel.history(limit=limit):
message_writer.writerow([
message.id,
message.author.id,
int(time.mktime(message.created_at.timetuple())),
int(time.mktime(message.edited_at.timetuple())) if message.edited_at else 0,
message.content
])
if message.attachments:
for attachment in message.attachments:
attachment_writer.writerow([
message.id,
attachment.content_type,
attachment.size,
attachment.filename,
attachment.url,
attachment.is_spoiler()
])
if message.reactions:
for reaction in message.reactions:
reaction_writer.writerow([
message.id,
reaction.emoji if isinstance(reaction.emoji, str) else reaction.emoji.name,
reaction.count,
False if isinstance(reaction.emoji, str) else reaction.emoji.animated
])
if message.embeds:
for embed in message.embeds:
embed_writer.writerow([
message.id,
embed.title if embed.title else None,
embed.description if embed.description else None,
embed.footer.text if embed.footer else None,
embed.image.url if embed.image else None,
embed.thumbnail.url if embed.thumbnail else None,
embed.video.url if embed.video else None,
embed.author.name if embed.author else None,
len(embed.fields),
embed.color if embed.color else None
])
async def main():
client = await request_client()
try:
guild = await request_guild(client)
directory = (await request_directory()).joinpath(str(guild.id))
if directory.exists():
print(f'Directory for guild already exists ({directory.absolute()})')
return
else:
directory.mkdir()
max_members = await request_max_members()
max_messages = await request_max_messages()
selected_channels = await request_channels(guild)
print('Starting dumping, this make take a while...')
with directory.joinpath('channels.csv').open('w+') as channels:
await dump_channels(guild, csv.writer(channels))
print('Dumped channels')
with directory.joinpath('members.csv').open('w+') as members:
await dump_members(guild, max_members, csv.writer(members))
print('Dumped members')
with directory.joinpath('roles.csv').open('w+') as roles:
await dump_roles(guild, csv.writer(roles))
print('Dumped roles')
for channel in selected_channels:
if isinstance(channel, discord.TextChannel):
dir = directory.joinpath(str(channel.id))
dir.mkdir()
with dir.joinpath('messages.csv').open('w+') as messages, dir.joinpath('attachments.csv').open('w+') as attachments, dir.joinpath('reactions.csv').open('w+') as reactions, dir.joinpath('embeds.csv').open('w+') as embeds:
await dump_messages(channel, max_messages, csv.writer(messages), csv.writer(attachments), csv.writer(reactions), csv.writer(embeds))
print(f'Dumped channel {channel.id} (#{channel.name})')
print('Finished dumping')
finally:
await client.close()
if __name__ == '__main__':
try:
asyncio.run(main())
except RuntimeError:
pass