klemek
143 lines
5.3 KiB
Python
143 lines
5.3 KiB
Python
from typing import Union, List
|
|
import os
|
|
import discord
|
|
import json
|
|
|
|
LOG_DIR = "logs"
|
|
|
|
if not os.path.exists(LOG_DIR):
|
|
os.mkdir(LOG_DIR)
|
|
|
|
|
|
CHUNK_SIZE = 1000
|
|
|
|
|
|
class MessageLog:
|
|
def __init__(self, message: Union[discord.Message, dict]):
|
|
if isinstance(message, discord.Message):
|
|
self.id = message.id
|
|
self.created_at = message.created_at
|
|
self.edited_at = message.edited_at
|
|
self.author = message.author
|
|
self.pinned = message.pinned
|
|
self.mention_everyone = message.mention_everyone
|
|
self.tts = message.tts
|
|
self.reference = message.reference.id
|
|
self.content = message.content
|
|
self.mentions = message.raw_mentions
|
|
self.role_mentions = message.raw_role_mentions
|
|
self.channel_mentions = message.raw_channel_mentions
|
|
self.reactions = {}
|
|
elif isinstance(message, dict):
|
|
self.id = message["id"]
|
|
self.created_at = message["created_at"]
|
|
self.edited_at = message["edited_at"]
|
|
self.author = message["author"]
|
|
self.pinned = message["pinned"]
|
|
self.mention_everyone = message["mention_everyone"]
|
|
self.tts = message["tts"]
|
|
self.reference = message["reference.id"]
|
|
self.content = message["content"]
|
|
self.mentions = message["raw_mentions"]
|
|
self.role_mentions = message["raw_role_mentions"]
|
|
self.channel_mentions = message["raw_channel_mentions"]
|
|
self.reactions = message["reactions"]
|
|
|
|
async def load(self, message: discord.Message):
|
|
for reaction in message.reactions:
|
|
self.reactions[str(reaction)] = []
|
|
async for user in reaction.users():
|
|
self.reactions[str(reaction)] += user.id
|
|
|
|
def dict(self):
|
|
return self.__dict__
|
|
|
|
|
|
class ChannelLog:
|
|
def __init__(self, channel: Union[discord.TextChannel, dict]):
|
|
if isinstance(channel, discord.TextChannel):
|
|
self.id = channel.id
|
|
self.name = channel.name
|
|
self.last_message_id = None
|
|
self.messages = []
|
|
elif isinstance(channel, dict):
|
|
self.id = channel["id"]
|
|
self.name = channel["name"]
|
|
self.last_message_id = channel["last_message_id"]
|
|
self.messages = [MessageLog(message) for message in channel["messages"]]
|
|
|
|
async def load(self, channel: discord.TextChannel):
|
|
self.name = channel.name
|
|
if self.last_message_id is not None: # append
|
|
while self.last_message_id != channel.last_message_id:
|
|
async for message in channel.history(
|
|
limit=CHUNK_SIZE, after=self.last_message_id, oldest_first=True
|
|
):
|
|
self.last_message_id = message.id
|
|
m = MessageLog(message)
|
|
await m.load(message)
|
|
self.messages.insert(0, m)
|
|
yield len(self.messages), False
|
|
else: # first load
|
|
last_message_id = None
|
|
done = 0
|
|
while done >= CHUNK_SIZE or last_message_id is None:
|
|
done = 0
|
|
async for message in channel.history(
|
|
limit=CHUNK_SIZE, after=self.last_message_id, oldest_first=False
|
|
):
|
|
done += 1
|
|
last_message_id = message.id
|
|
m = MessageLog(message)
|
|
await m.load(message)
|
|
self.messages += [m]
|
|
yield len(self.messages), False
|
|
self.last_message_id == channel.last_message_id
|
|
yield len(self.messages), True
|
|
|
|
def dict(self):
|
|
tmp = self.__dict__
|
|
tmp["messages"] = [message.dict() for message in self.messages]
|
|
return tmp
|
|
|
|
|
|
class GuildLogs:
|
|
def __init__(self, guild: discord.Guild):
|
|
self.guild = guild
|
|
self.log_file = os.path.join(LOG_DIR, f"{guild}.logz")
|
|
self.channels = {}
|
|
|
|
def dict(self):
|
|
return {id: self.channels[id].dict() for id in self.channels}
|
|
|
|
async def load(self, target_channels: List[discord.TextChannel] = []):
|
|
# read logs
|
|
if os.path.exists(self.log_file):
|
|
channels = {}
|
|
with open(self.log_file, mode="r") as f:
|
|
channels = json.loads(f.readline().strip())
|
|
self.channels = {id: ChannelLog(channels[id]) for id in channels}
|
|
# load channels
|
|
if len(target_channels) == 0:
|
|
target_channels = self.guild.text_channels
|
|
loading_new = False
|
|
total_msg = 0
|
|
total_chan = 0
|
|
for channel in target_channels:
|
|
if channel.id not in self.channels:
|
|
loading_new = True
|
|
self.channels[channel.id] = ChannelLog(channel)
|
|
async for count, done in self.channels[channel.id].load(channel):
|
|
yield (
|
|
total_msg + count,
|
|
total_chan + (1 if done else 0),
|
|
loading_new,
|
|
False,
|
|
)
|
|
total_msg += len(self.channels[channel.id].messages)
|
|
total_chan += 1
|
|
yield total_msg, total_chan, loading_new, True
|
|
# write logs
|
|
with open(self.log_file, mode="w") as f:
|
|
f.write(json.dump(self.dict()))
|