Skip to content

Commit b7a4a87

Browse files
hypergonialhypergonial
committed
Refactor userlog
1 parent a9f3987 commit b7a4a87

File tree

5 files changed

+182
-74
lines changed

5 files changed

+182
-74
lines changed

‎extensions/starboard.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,8 @@ def create_starboard_payload(
9898

9999
if attachments:
100100
embed.add_field(
101-
"Attachments", "\n".join([f"[{attachment.filename[:100]}]({attachment.url})" for attachment in attachments])
101+
"Attachments",
102+
"\n".join([f"[{attachment.filename[:100]}]({attachment.url})" for attachment in attachments][:5]),
102103
)
103104

104105
if message.referenced_message:

‎extensions/userlog.py‎

Lines changed: 34 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -255,17 +255,20 @@ async def unfreeze_logging(guild_id: int) -> None:
255255

256256
userlog.d.actions["unfreeze_logging"] = unfreeze_logging
257257

258-
258+
# TODO: refactor this to better take advantage of the new audit log cache
259259
async def find_auditlog_data(
260-
event: hikari.Event, *, event_type: hikari.AuditLogEventType, user_id: t.Optional[int] = None
260+
guild: hikari.SnowflakeishOr[hikari.PartialGuild],
261+
*,
262+
event_type: hikari.AuditLogEventType,
263+
user_id: t.Optional[int] = None,
261264
) -> t.Optional[hikari.AuditLogEntry]:
262265
"""Find a recently sent audit log entry that matches criteria.
263266
264267
Parameters
265268
----------
266-
event : hikari.GuildEvent
269+
event : hikari.Event
267270
The event that triggered this search.
268-
type : hikari.AuditLogEventType
271+
event_type : hikari.AuditLogEventType
269272
The type of audit log entry to look for.
270273
user_id : Optional[int], optional
271274
The user affected by this audit log, if any. By default hikari.UNDEFINED
@@ -284,41 +287,12 @@ async def find_auditlog_data(
284287
# Stuff that is observed to just take too goddamn long to appear in AuditLogs
285288
takes_an_obscene_amount_of_time = [hikari.AuditLogEventType.MESSAGE_BULK_DELETE]
286289

287-
guild = userlog.app.cache.get_guild(event.guild_id) # type: ignore
288-
sleep_time = 15.0 if event_type not in takes_an_obscene_amount_of_time else 30.0
289-
await asyncio.sleep(sleep_time) # Wait for auditlog to hopefully fill in
290-
291-
if not guild:
292-
raise ValueError("Cannot find guild to parse auditlogs for.")
293-
294-
me = userlog.app.cache.get_member(guild, userlog.app.user_id)
295-
296-
if me is None:
297-
return
298-
299-
perms = lightbulb.utils.permissions_for(me)
300-
if not (perms & hikari.Permissions.VIEW_AUDIT_LOG):
301-
# Do not attempt to fetch audit log if bot has no perms
302-
return
303-
304-
try:
305-
count = 0
306-
async for log in userlog.app.rest.fetch_audit_log(guild, event_type=event_type):
307-
for entry in log.entries.values():
308-
# We do not want to return entries older than 15 seconds
309-
if (helpers.utcnow() - entry.id.created_at).total_seconds() > 30 or count > 5:
310-
return
311-
312-
if user_id and user_id == entry.target_id:
313-
return entry
314-
elif user_id:
315-
count += 1
316-
continue
317-
else:
318-
return entry
290+
sleep_time = 5.0 if event_type not in takes_an_obscene_amount_of_time else 30.0
291+
await asyncio.sleep(sleep_time) # Wait for auditlog event to hopefully arrive
319292

320-
except (hikari.ForbiddenError, hikari.HTTPError, asyncio.TimeoutError):
321-
return
293+
return userlog.app.audit_log_cache.get_first_by(
294+
guild, event_type, lambda e: e.target_id == user_id if user_id else True
295+
)
322296

323297

324298
async def get_perms_diff(old_role: hikari.Role, role: hikari.Role) -> t.Optional[str]:
@@ -446,7 +420,7 @@ async def message_delete(plugin: SnedPlugin, event: hikari.GuildMessageDeleteEve
446420
contents = create_log_content(event.old_message)
447421

448422
entry = await find_auditlog_data(
449-
event, event_type=hikari.AuditLogEventType.MESSAGE_DELETE, user_id=event.old_message.author.id
423+
event.guild_id, event_type=hikari.AuditLogEventType.MESSAGE_DELETE, user_id=event.old_message.author.id
450424
)
451425

452426
if entry:
@@ -505,7 +479,7 @@ async def message_update(plugin: SnedPlugin, event: hikari.GuildMessageUpdateEve
505479
async def bulk_message_delete(plugin: SnedPlugin, event: hikari.GuildBulkMessageDeleteEvent) -> None:
506480

507481
moderator = "Discord"
508-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.MESSAGE_BULK_DELETE)
482+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.MESSAGE_BULK_DELETE)
509483
if entry:
510484
assert entry.user_id is not None
511485
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id)
@@ -525,7 +499,7 @@ async def bulk_message_delete(plugin: SnedPlugin, event: hikari.GuildBulkMessage
525499
@userlog.listener(hikari.RoleDeleteEvent, bind=True)
526500
async def role_delete(plugin: SnedPlugin, event: hikari.RoleDeleteEvent) -> None:
527501

528-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.ROLE_DELETE)
502+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.ROLE_DELETE)
529503
if entry and event.old_role:
530504
assert entry.user_id is not None
531505
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id)
@@ -540,7 +514,7 @@ async def role_delete(plugin: SnedPlugin, event: hikari.RoleDeleteEvent) -> None
540514
@userlog.listener(hikari.RoleCreateEvent, bind=True)
541515
async def role_create(plugin: SnedPlugin, event: hikari.RoleCreateEvent) -> None:
542516

543-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.ROLE_CREATE)
517+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.ROLE_CREATE)
544518
if entry and event.role:
545519
assert entry.user_id is not None
546520
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id)
@@ -555,7 +529,7 @@ async def role_create(plugin: SnedPlugin, event: hikari.RoleCreateEvent) -> None
555529
@userlog.listener(hikari.RoleUpdateEvent, bind=True)
556530
async def role_update(plugin: SnedPlugin, event: hikari.RoleUpdateEvent) -> None:
557531

558-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.ROLE_UPDATE)
532+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.ROLE_UPDATE)
559533
if entry and event.old_role:
560534
assert entry.user_id
561535
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id)
@@ -586,7 +560,7 @@ async def role_update(plugin: SnedPlugin, event: hikari.RoleUpdateEvent) -> None
586560
@userlog.listener(hikari.GuildChannelDeleteEvent, bind=True)
587561
async def channel_delete(plugin: SnedPlugin, event: hikari.GuildChannelDeleteEvent) -> None:
588562

589-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.CHANNEL_DELETE)
563+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.CHANNEL_DELETE)
590564
if entry and event.channel:
591565
assert entry.user_id is not None
592566
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id)
@@ -601,7 +575,7 @@ async def channel_delete(plugin: SnedPlugin, event: hikari.GuildChannelDeleteEve
601575
@userlog.listener(hikari.GuildChannelCreateEvent, bind=True)
602576
async def channel_create(plugin: SnedPlugin, event: hikari.GuildChannelCreateEvent) -> None:
603577

604-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.CHANNEL_CREATE)
578+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.CHANNEL_CREATE)
605579
if entry and event.channel:
606580
assert entry.user_id is not None
607581
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id)
@@ -616,7 +590,7 @@ async def channel_create(plugin: SnedPlugin, event: hikari.GuildChannelCreateEve
616590
@userlog.listener(hikari.GuildChannelUpdateEvent, bind=True)
617591
async def channel_update(plugin: SnedPlugin, event: hikari.GuildChannelUpdateEvent) -> None:
618592

619-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.CHANNEL_UPDATE)
593+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.CHANNEL_UPDATE)
620594

621595
if entry and event.old_channel:
622596
assert entry.user_id is not None
@@ -663,7 +637,7 @@ async def channel_update(plugin: SnedPlugin, event: hikari.GuildChannelUpdateEve
663637
@userlog.listener(hikari.GuildUpdateEvent, bind=True)
664638
async def guild_update(plugin: SnedPlugin, event: hikari.GuildUpdateEvent) -> None:
665639

666-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.GUILD_UPDATE)
640+
entry = await find_auditlog_data(event.guild_id, event_type=hikari.AuditLogEventType.GUILD_UPDATE)
667641

668642
if event.old_guild:
669643
if entry:
@@ -722,7 +696,7 @@ async def guild_update(plugin: SnedPlugin, event: hikari.GuildUpdateEvent) -> No
722696
async def member_ban_remove(plugin: SnedPlugin, event: hikari.BanDeleteEvent) -> None:
723697

724698
entry = await find_auditlog_data(
725-
event, event_type=hikari.AuditLogEventType.MEMBER_BAN_REMOVE, user_id=event.user.id
699+
event.guild_id, event_type=hikari.AuditLogEventType.MEMBER_BAN_REMOVE, user_id=event.user.id
726700
)
727701
if entry:
728702
assert entry.user_id is not None
@@ -757,7 +731,9 @@ async def member_ban_remove(plugin: SnedPlugin, event: hikari.BanDeleteEvent) ->
757731
@userlog.listener(hikari.BanCreateEvent, bind=True)
758732
async def member_ban_add(plugin: SnedPlugin, event: hikari.BanCreateEvent) -> None:
759733

760-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.MEMBER_BAN_ADD, user_id=event.user.id)
734+
entry = await find_auditlog_data(
735+
event.guild_id, event_type=hikari.AuditLogEventType.MEMBER_BAN_ADD, user_id=event.user.id
736+
)
761737
if entry:
762738
assert entry.user_id is not None
763739
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id) if entry else "Unknown"
@@ -794,7 +770,9 @@ async def member_delete(plugin: SnedPlugin, event: hikari.MemberDeleteEvent) ->
794770
if event.user_id == plugin.app.user_id:
795771
return # RIP
796772

797-
entry = await find_auditlog_data(event, event_type=hikari.AuditLogEventType.MEMBER_KICK, user_id=event.user.id)
773+
entry = await find_auditlog_data(
774+
event.guild_id, event_type=hikari.AuditLogEventType.MEMBER_KICK, user_id=event.user.id
775+
)
798776

799777
if entry: # This is a kick
800778
assert entry.user_id is not None
@@ -864,7 +842,7 @@ async def member_update(plugin: SnedPlugin, event: hikari.MemberUpdateEvent) ->
864842
comms_disabled_until = member.communication_disabled_until()
865843

866844
entry = await find_auditlog_data(
867-
event, event_type=hikari.AuditLogEventType.MEMBER_UPDATE, user_id=event.user.id
845+
event.guild_id, event_type=hikari.AuditLogEventType.MEMBER_UPDATE, user_id=event.user.id
868846
)
869847
if not entry:
870848
return
@@ -943,10 +921,13 @@ async def member_update(plugin: SnedPlugin, event: hikari.MemberUpdateEvent) ->
943921
return
944922

945923
entry = await find_auditlog_data(
946-
event, event_type=hikari.AuditLogEventType.MEMBER_ROLE_UPDATE, user_id=event.user.id
924+
event.guild_id, event_type=hikari.AuditLogEventType.MEMBER_ROLE_UPDATE, user_id=event.user.id
947925
)
948926

949-
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id) if entry and entry.user_id else "Unknown"
927+
moderator = plugin.app.cache.get_member(event.guild_id, entry.user_id) if entry and entry.user_id else None
928+
929+
if moderator is None:
930+
return
950931

951932
if entry and entry.user_id == plugin.app.user_id:
952933
# Attempt to find the moderator in reason if sent by us

‎models/audit_log.py‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import typing as t
2+
3+
import hikari
4+
5+
if t.TYPE_CHECKING:
6+
from models.bot import SnedBot
7+
8+
import logging
9+
10+
logger = logging.getLogger(__name__)
11+
12+
13+
class AuditLogCache:
14+
"""A cache for audit log entries categorized by guild and entry type.
15+
16+
Parameters
17+
----------
18+
bot: SnedBot
19+
The bot instance.
20+
capacity: int
21+
The maximum number of entries to store per guild and event type. If the number
22+
of entries exceed this number, the oldest entries will be discarded.
23+
"""
24+
25+
def __init__(self, bot: SnedBot, capacity: int = 10) -> None:
26+
self._cache: t.Dict[hikari.Snowflake, t.Dict[hikari.AuditLogEventType, t.List[hikari.AuditLogEntry]]] = {}
27+
self._capacity = capacity
28+
self._bot = bot
29+
30+
async def start(self) -> None:
31+
"""Start the audit log cache listener."""
32+
self._bot.event_manager.subscribe(hikari.Event, self._listen)
33+
34+
async def stop(self) -> None:
35+
"""Stop the audit log cache listener."""
36+
self._bot.event_manager.unsubscribe(hikari.Event, self._listen)
37+
self._cache = {}
38+
39+
async def _listen(self, event: hikari.Event) -> None:
40+
"""Listen for audit log events."""
41+
raise NotImplementedError("AuditLogCache listener not implemented!")
42+
# TODO: do impl:
43+
# self.add(event.guild_id, event.entry)
44+
45+
def get(
46+
self, guild: hikari.SnowflakeishOr[hikari.PartialGuild], action_type: hikari.AuditLogEventType
47+
) -> t.List[hikari.AuditLogEntry]:
48+
"""Get all audit log entries for a guild and event type.
49+
50+
Parameters
51+
----------
52+
guild: hikari.SnowflakeishOr[hikari.PartialGuild]
53+
The guild or it's ID.
54+
action_type: hikari.AuditLogEventType
55+
The event type.
56+
57+
Returns
58+
-------
59+
List[hikari.AuditLogEntry]
60+
The audit log entries.
61+
"""
62+
return self._cache.get(hikari.Snowflake(guild), {}).get(action_type, [])
63+
64+
def get_first_by(
65+
self,
66+
guild: hikari.SnowflakeishOr[hikari.PartialGuild],
67+
action_type: hikari.AuditLogEventType,
68+
predicate: t.Callable[[hikari.AuditLogEntry], bool],
69+
) -> t.Optional[hikari.AuditLogEntry]:
70+
"""Get the first audit log entry that matches a predicate.
71+
72+
Parameters
73+
----------
74+
guild: hikari.SnowflakeishOr[hikari.PartialGuild]
75+
The guild or it's ID.
76+
action_type: hikari.AuditLogEventType
77+
The event type.
78+
predicate: Callable[[hikari.AuditLogEntry], bool]
79+
The predicate to match.
80+
81+
Returns
82+
-------
83+
Optional[hikari.AuditLogEntry]
84+
The first audit log entry that matches the predicate, or None if no entry matches.
85+
"""
86+
for entry in self.get(guild, action_type):
87+
if predicate(entry):
88+
return entry
89+
90+
return None
91+
92+
def add(self, guild: hikari.SnowflakeishOr[hikari.PartialGuild], entry: hikari.AuditLogEntry) -> None:
93+
"""Add a new audit log entry to the cache.
94+
95+
Parameters
96+
----------
97+
guild: hikari.SnowflakeishOr[hikari.PartialGuild]
98+
The guild or it's ID.
99+
entry: hikari.AuditLogEntry
100+
The audit log entry to add.
101+
"""
102+
if not isinstance(entry.action_type, hikari.AuditLogEventType):
103+
logger.warning(f"Unrecognized audit log entry type found: {entry.action_type}")
104+
return
105+
106+
guild_id = hikari.Snowflake(guild)
107+
108+
if guild_id not in self._cache:
109+
self._cache[guild_id] = {}
110+
111+
if entry.action_type not in self._cache[guild_id]:
112+
self._cache[guild_id][entry.action_type] = []
113+
114+
# Remove the oldest entry if the cache is full
115+
if len(self._cache[guild_id][entry.action_type]) >= self._capacity:
116+
self._cache[guild_id][entry.action_type].pop(0)
117+
118+
self._cache[guild_id][entry.action_type].append(entry)

‎models/bot.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
import utils.db_backup as db_backup
1414
from config import Config
15+
from models.audit_log import AuditLogCache
1516
from models.db import Database
1617
from models.errors import UserBlacklistedError
1718
from models.mod_actions import ModActions
@@ -110,6 +111,7 @@ def __init__(self, config: Config) -> None:
110111
self._user_id: t.Optional[hikari.Snowflake] = None
111112
self._perspective: t.Optional[kosu.Client] = None
112113
self._scheduler = scheduler.Scheduler(self)
114+
self._audit_log_cache: AuditLogCache = AuditLogCache(self)
113115
self._initial_guilds: t.List[hikari.Snowflake] = []
114116

115117
self.check(is_not_blacklisted)
@@ -176,6 +178,11 @@ def mod(self) -> ModActions:
176178
"""The moderation actions instance of the bot. Handles moderation of users and contains useful methods for such purposes."""
177179
return self._mod
178180

181+
@property
182+
def audit_log_cache(self) -> AuditLogCache:
183+
"""The audit log cache instance of the bot."""
184+
return self._audit_log_cache
185+
179186
@property
180187
def is_started(self) -> bool:
181188
"""Boolean indicating if the bot has started up or not."""
@@ -242,6 +249,7 @@ async def on_starting(self, event: hikari.StartingEvent) -> None:
242249
# Start scheduler, DB cache
243250
await self.db_cache.start()
244251
self.scheduler.start()
252+
await self._audit_log_cache.start()
245253

246254
if perspective_api_key := os.getenv("PERSPECTIVE_API_KEY"):
247255
self._perspective = kosu.Client(perspective_api_key, do_not_store=True)

0 commit comments

Comments
 (0)