Rocketsoup 4 лет назад
Родитель
Сommit
9e8be49748
6 измененных файлов: 83 добавлений и 36 удалений
  1. 2
    1
      .gitignore
  2. 35
    5
      cogs/base.py
  3. 8
    9
      cogs/config.py
  4. 8
    1
      cogs/general.py
  5. 2
    1
      cogs/joinraid.py
  6. 28
    19
      storage.py

+ 2
- 1
.gitignore Просмотреть файл

117
 # Mac!!
117
 # Mac!!
118
 .DS_Store
118
 .DS_Store
119
 
119
 
120
-# Contains secrets so exclude
120
+# Rocketbot stuff
121
 /config.py
121
 /config.py
122
 rocketbot.db
122
 rocketbot.db
123
+/state/

+ 35
- 5
cogs/base.py Просмотреть файл

1
-from discord import Guild
1
+from discord import Guild, Message, TextChannel
2
 from discord.ext import commands
2
 from discord.ext import commands
3
 
3
 
4
-from storage import Storage
4
+from storage import StateKey, Storage
5
 
5
 
6
 class BaseCog(commands.Cog):
6
 class BaseCog(commands.Cog):
7
 	def __init__(self, bot):
7
 	def __init__(self, bot):
8
 		self.bot = bot
8
 		self.bot = bot
9
 
9
 
10
+	@staticmethod
10
 	def save_setting(guild: Guild, name: str, value):
11
 	def save_setting(guild: Guild, name: str, value):
12
+		"""
13
+		Saves one value to a guild's persisted state. The given value must
14
+		be a JSON-encodable type.
15
+		"""
16
+		if value is not None and not isinstance(value, (bool, int, float, str, list, dict)):
17
+			raise Exception(f'value for key {name} is not supported JSON type! {type(value)}')
11
 		state: dict = Storage.state_for_guild(guild)
18
 		state: dict = Storage.state_for_guild(guild)
12
-		state[name] = value
19
+		if value is None:
20
+			del state[name]
21
+		else:
22
+			state[name] = value
13
 		Storage.save_guild_state(guild)
23
 		Storage.save_guild_state(guild)
14
 
24
 
15
-	def warn(guild: Guild):
25
+	@staticmethod
26
+	async def warn(guild: Guild, message: str) -> bool:
27
+		"""
28
+		Sends a warning message to the configured warning channel for the
29
+		given guild. If no warning channel is configured no action is taken.
30
+		Returns True if the message was sent successfully or False if the
31
+		warning channel could not be found.
32
+		"""
16
 		state: dict = Storage.state_for_guild(guild)
33
 		state: dict = Storage.state_for_guild(guild)
17
-		
34
+		channel_id = state.get(StateKey.WARNING_CHANNEL_ID)
35
+		if channel_id is None:
36
+			guild_trace(guild, 'No warning channel set! No warning issued.')
37
+			return False
38
+		channel: TextChannel = guild.get_channel(channel_id)
39
+		if channel is None:
40
+			guild_trace(guild, 'Configured warning channel no longer exists!')
41
+			return False
42
+		message: Message = await channel.send(message)
43
+		return message is not None
44
+
45
+	@staticmethod
46
+	def guild_trace(guild: Guild, message: str) -> None:
47
+		print(f'[guild {guild.id}|{guild.name}] {message}')

+ 8
- 9
cogs/config.py Просмотреть файл

1
 from discord import Guild, TextChannel
1
 from discord import Guild, TextChannel
2
 from discord.ext import commands
2
 from discord.ext import commands
3
-from storage import Storage
3
+from storage import StateKey, Storage
4
+from cogs.base import BaseCog
4
 
5
 
5
-class ConfigCog(commands.Cog):
6
+class ConfigCog(BaseCog):
6
 	"""
7
 	"""
7
 	Cog for handling general bot configuration.
8
 	Cog for handling general bot configuration.
8
 	"""
9
 	"""
20
 			await context.send_help()
21
 			await context.send_help()
21
 
22
 
22
 	@config.command(
23
 	@config.command(
23
-		name='setoutputchannel',
24
-		brief='Sets the channel where mod communications occur',
24
+		name='setwarningchannel',
25
+		brief='Sets the channel where mod warnings are posted',
25
 		description='Run this command in the channel where bot messages ' +
26
 		description='Run this command in the channel where bot messages ' +
26
 			'intended for server moderators should be sent. Other bot messages ' +
27
 			'intended for server moderators should be sent. Other bot messages ' +
27
 			'may still be posted in the channel a command was invoked in. If ' +
28
 			'may still be posted in the channel a command was invoked in. If ' +
28
 			'no output channel is set, mod-related messages will not be posted!',
29
 			'no output channel is set, mod-related messages will not be posted!',
29
 	)
30
 	)
30
-	async def config_setoutputchannel(self, context: commands.Context):
31
+	async def config_setwarningchannel(self, context: commands.Context):
31
 		'Command handler'
32
 		'Command handler'
32
 		guild: Guild = context.guild
33
 		guild: Guild = context.guild
33
 		channel: TextChannel = context.channel
34
 		channel: TextChannel = context.channel
34
-		state: dict = Storage.state_for_guild(context.guild)
35
-		state['output_channel_id'] = context.channel.id
36
-		Storage.save_guild_state(context.guild)
37
-		await context.message.reply(f'Output channel set to {channel.name}', mention_author=False)
35
+		self.save_setting(guild, StateKey.WARNING_CHANNEL_ID, context.channel.id)
36
+		await context.message.reply(f'Warning channel set to {channel.name}', mention_author=False)

+ 8
- 1
cogs/general.py Просмотреть файл

1
 from discord.ext import commands
1
 from discord.ext import commands
2
+from cogs.base import BaseCog
2
 
3
 
3
-class GeneralCog(commands.Cog):
4
+class GeneralCog(BaseCog):
4
 	def __init__(self, bot: commands.Bot):
5
 	def __init__(self, bot: commands.Bot):
5
 		self.bot = bot
6
 		self.bot = bot
6
 		self.is_connected = False
7
 		self.is_connected = False
15
 	async def on_ready(self):
16
 	async def on_ready(self):
16
 		print('on_ready')
17
 		print('on_ready')
17
 		self.is_ready = True
18
 		self.is_ready = True
19
+
20
+	@commands.command()
21
+	@commands.has_permissions(ban_members=True)
22
+	@commands.guild_only()
23
+	async def testwarn(self, context):
24
+		await self.warn(context.guild, 'Test warning')

+ 2
- 1
cogs/joinraid.py Просмотреть файл

1
 from discord import Guild, Intents, Member, Message, PartialEmoji, RawReactionActionEvent
1
 from discord import Guild, Intents, Member, Message, PartialEmoji, RawReactionActionEvent
2
 from discord.ext import commands
2
 from discord.ext import commands
3
 from storage import Storage
3
 from storage import Storage
4
+from cogs.base import BaseCog
4
 
5
 
5
 class JoinRecord:
6
 class JoinRecord:
6
     """
7
     """
149
         self.phase = RaidPhase.ENDED
150
         self.phase = RaidPhase.ENDED
150
         return bans
151
         return bans
151
 
152
 
152
-class JoinRaidCog(commands.Cog):
153
+class JoinRaidCog(BaseCog):
153
 	"""
154
 	"""
154
 	Cog for monitoring member joins and detecting potential bot raids.
155
 	Cog for monitoring member joins and detecting potential bot raids.
155
 	"""
156
 	"""

+ 28
- 19
storage.py Просмотреть файл

5
 
5
 
6
 from config import CONFIG
6
 from config import CONFIG
7
 
7
 
8
+class StateKey:
9
+	WARNING_CHANNEL_ID = 'warning_channel_id'
10
+
8
 class Storage:
11
 class Storage:
9
 	"""
12
 	"""
10
 	Static class for managing persisted bot state.
13
 	Static class for managing persisted bot state.
13
 	# discord.Guild.id -> dict
16
 	# discord.Guild.id -> dict
14
 	__guild_id_to_state = {}
17
 	__guild_id_to_state = {}
15
 
18
 
16
-	def state_for_guild(guild: Guild) -> dict:
19
+	@classmethod
20
+	def state_for_guild(cls, guild: Guild) -> dict:
17
 		"""
21
 		"""
18
 		Returns the state for the given guild, loading from disk if necessary.
22
 		Returns the state for the given guild, loading from disk if necessary.
19
 		Always returns a dict.
23
 		Always returns a dict.
20
 		"""
24
 		"""
21
-		state: dict = __guild_id_to_state[guild.id]
25
+		state: dict = cls.__guild_id_to_state.get(guild.id)
22
 		if state is not None:
26
 		if state is not None:
23
 			# Already in memory
27
 			# Already in memory
24
 			return state
28
 			return state
25
 		# Load from disk if possible
29
 		# Load from disk if possible
26
-		__trace(f'No loaded state for guild {guild.id}. Attempting to load from disk.')
27
-		state = __load_guild_state(guild)
28
-		__guild_id_to_state[guild.id] = state
30
+		cls.__trace(f'No loaded state for guild {guild.id}. Attempting to load from disk.')
31
+		state = cls.__load_guild_state(guild)
32
+		cls.__guild_id_to_state[guild.id] = state
29
 		return state
33
 		return state
30
 
34
 
31
-	def save_guild_state(guild: Guild) -> NoneType:
35
+	@classmethod
36
+	def save_guild_state(cls, guild: Guild) -> None:
32
 		"""
37
 		"""
33
 		Saves state for the given guild to disk, if any exists.
38
 		Saves state for the given guild to disk, if any exists.
34
 		"""
39
 		"""
35
-		state: dict = __guild_id_to_state[guild.id]
40
+		state: dict = cls.__guild_id_to_state.get(guild.id)
36
 		if state is None:
41
 		if state is None:
37
 			# Nothing to save
42
 			# Nothing to save
38
 			return
43
 			return
39
-		__save_guild_state(guild, state)
44
+		cls.__save_guild_state(guild, state)
40
 
45
 
41
-	def __save_guild_state(guild: Guild, state: dict) -> NoneType:
46
+	@classmethod
47
+	def __save_guild_state(cls, guild: Guild, state: dict) -> None:
42
 		"""
48
 		"""
43
 		Saves state for a guild to a JSON file on disk.
49
 		Saves state for a guild to a JSON file on disk.
44
 		"""
50
 		"""
45
-		path: str = __guild_path(guild)
46
-		__trace('Saving state for guild {guild.id} to {path}')
51
+		path: str = cls.__guild_path(guild)
52
+		cls.__trace('Saving state for guild {guild.id} to {path}')
47
 		with open(path, 'w') as file:
53
 		with open(path, 'w') as file:
48
 			# Pretty printing to make more legible for debugging
54
 			# Pretty printing to make more legible for debugging
49
 			# Sorting keys to help with diffs
55
 			# Sorting keys to help with diffs
50
 			json.dump(state, file, indent='\t', sort_keys=True)
56
 			json.dump(state, file, indent='\t', sort_keys=True)
51
-		__trace('State saved')
57
+		cls.__trace('State saved')
52
 
58
 
53
-	def __load_guild_state(guild: Guild) -> dict:
59
+	@classmethod
60
+	def __load_guild_state(cls, guild: Guild) -> dict:
54
 		"""
61
 		"""
55
 		Loads state for a guild from a JSON file on disk.
62
 		Loads state for a guild from a JSON file on disk.
56
 		"""
63
 		"""
57
-		path: str = __guild_path(guild)
64
+		path: str = cls.__guild_path(guild)
58
 		if not exists(path):
65
 		if not exists(path):
59
-			__trace(f'No state on disk for guild {guild.id}. Returning {{}}.')
66
+			cls.__trace(f'No state on disk for guild {guild.id}. Returning {{}}.')
60
 			return {}
67
 			return {}
61
-		__trace('Loading state from disk for guild {guild.id}')
68
+		cls.__trace(f'Loading state from disk for guild {guild.id}')
62
 		with open(path, 'r') as file:
69
 		with open(path, 'r') as file:
63
 			state = json.load(file)
70
 			state = json.load(file)
64
-		__trace('State loaded')
71
+		cls.__trace('State loaded')
65
 		return state
72
 		return state
66
 
73
 
67
-	def __guild_path(guild: Guild) -> str:
74
+	@classmethod
75
+	def __guild_path(cls, guild: Guild) -> str:
68
 		"""
76
 		"""
69
 		Returns the JSON file path where guild state should be written.
77
 		Returns the JSON file path where guild state should be written.
70
 		"""
78
 		"""
72
 		path: str = config_value if config_value.endswith('/') else f'{config_value}/'
80
 		path: str = config_value if config_value.endswith('/') else f'{config_value}/'
73
 		return f'{path}guild_{guild.id}.json'
81
 		return f'{path}guild_{guild.id}.json'
74
 
82
 
75
-	def __trace(message: str) -> NoneType:
83
+	@classmethod
84
+	def __trace(cls, message: str) -> None:
76
 		print(f'{Storage.__name__}: {message}')
85
 		print(f'{Storage.__name__}: {message}')

Загрузка…
Отмена
Сохранить