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
|
import discord
from discord import app_commands
from discord.ext import commands
import aiohttp
import random
import string
import os
from src.utils.db import get_db
from models import Message
from config import BOT_COLOR, LOG
class Archive(commands.Cog):
def __init__(self, bot):
self.bot = bot
async def download_attachments(self, attachments) -> list:
"""Download attachments and return a list of their paths."""
paths = []
for attachment in attachments:
async with aiohttp.ClientSession() as session:
async with session.get(attachment.url) as response:
# Check if the request was successful
if response.status != 200:
LOG.warn(
f"Failed to download attachment: {attachment.url}"
)
continue
# Check for content type
content_type = response.headers.get("Content-Type")
if not content_type:
LOG.warn(
f"Failed to get content type for: {attachment.url}"
)
continue
# Create a randomized filename
file_extension = content_type.split("/")[-1]
filename = (
"".join(
random.choice(string.ascii_letters)
for i in range(10)
)
+ f".{file_extension}"
)
# Save the attachment
with open(f"images/{filename}", "wb") as file:
file.write(await response.read())
# Add the path to the attachments list
paths.append(f"images/{filename}")
return paths
@app_commands.command()
async def archive(
self,
interaction: discord.Interaction,
channel: discord.TextChannel,
amount: int,
):
"""Archive a channel's messages."""
# Ensure valid channel permissions
if not channel.permissions_for(
interaction.guild.me
).read_message_history:
return await interaction.response.send_message(
"I do not have permission to read message history in that"
" channel.",
ephemeral=True,
)
# Ensure valid amount
if amount < 1:
return await interaction.response.send_message(
"You must provide a number greater than 0.",
ephemeral=True,
)
embed = discord.Embed(
title="Archive Beginning",
description=(
f"Archiving {amount} messages from {channel.mention}. I will"
" DM you a message once complete, make sure to allow messages"
" from me."
),
color=BOT_COLOR,
)
await interaction.response.send_message(embed=embed, ephemeral=True)
db = next(get_db())
count = 0
messages = channel.history(limit=amount)
async for message in messages:
count += 1
if not os.path.exists("images"):
os.makedirs("images")
attachments = await self.download_attachments(message.attachments)
db_message = Message(
timestamp=message.created_at.isoformat(),
message_id=message.id,
author_id=message.author.id,
channel_id=channel.id,
stickers=[sticker.name for sticker in message.stickers],
role_mentions=[role.id for role in message.role_mentions],
mention_everyone=message.mention_everyone,
mentions=[mention.id for mention in message.mentions],
attachments=attachments,
content=message.content,
)
db.add(db_message)
if count > 500:
db.commit()
count = 0
db.commit()
embed = discord.Embed(
title="Archive Complete",
description=f"Archived {amount} messages from {channel.mention}.",
color=BOT_COLOR,
)
try:
await interaction.user.send(embed=embed)
except discord.Forbidden:
await channel.send(
f"{interaction.user.mention} I have completed the archive, but"
" was unable to DM you the final message. Please check your"
" DM settings to receive future messages."
)
async def setup(bot):
await bot.add_cog(Archive(bot))
|
