aboutsummaryrefslogtreecommitdiffstats
path: root/arthur
diff options
context:
space:
mode:
Diffstat (limited to 'arthur')
-rw-r--r--arthur/__init__.py12
-rw-r--r--arthur/__main__.py37
-rw-r--r--arthur/apis/cloudflare/zones.py20
-rw-r--r--arthur/bot.py50
-rw-r--r--arthur/config.py3
-rw-r--r--arthur/extensions.py15
-rw-r--r--arthur/exts/__init__.py1
-rw-r--r--arthur/exts/cloudflare/__init__.py1
-rw-r--r--arthur/exts/cloudflare/zones.py16
-rw-r--r--arthur/exts/error_handler/__init__.py1
-rw-r--r--arthur/exts/error_handler/error_handler.py4
-rw-r--r--arthur/exts/fun/__init__.py1
-rw-r--r--arthur/exts/fun/ed.py8
-rw-r--r--arthur/exts/kubernetes/__init__.py1
-rw-r--r--arthur/exts/kubernetes/certificates.py4
-rw-r--r--arthur/exts/kubernetes/deployments.py8
-rw-r--r--arthur/exts/kubernetes/jobs.py6
-rw-r--r--arthur/exts/kubernetes/nodes.py4
18 files changed, 103 insertions, 89 deletions
diff --git a/arthur/__init__.py b/arthur/__init__.py
index a7efb0b..0d837f6 100644
--- a/arthur/__init__.py
+++ b/arthur/__init__.py
@@ -1,7 +1,19 @@
"""King Arthur is Python Discord's DevOps utility bot."""
+import asyncio
+import os
from functools import partial
+from typing import TYPE_CHECKING
import loguru
+if TYPE_CHECKING:
+ from arthur.bot import KingArthur
+
logger = loguru.logger.opt(colors=True)
logger.opt = partial(logger.opt, colors=True)
+
+# On Windows, the selector event loop is required for aiodns.
+if os.name == "nt":
+ asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+instance: "KingArthur" = None # Global Bot instance.
diff --git a/arthur/__main__.py b/arthur/__main__.py
index 7f95210..bc09f09 100644
--- a/arthur/__main__.py
+++ b/arthur/__main__.py
@@ -1,16 +1,37 @@
"""Entrypoint for King Arthur."""
-from arthur import logger
+import asyncio
+
+import aiohttp
+import discord
+from discord.ext import commands
+
+import arthur
from arthur.bot import KingArthur
from arthur.config import CONFIG
[email protected]()
-def start() -> None:
- """Entrypoint for King Arthur."""
- arthur = KingArthur()
+async def main() -> None:
+ """Entry async method for starting the bot."""
+ intents = discord.Intents.default()
+ intents.message_content = True
+ intents.dm_typing = False
+ intents.dm_reactions = False
+ intents.invites = False
+ intents.webhooks = False
+ intents.integrations = False
- arthur.run(CONFIG.token)
+ async with aiohttp.ClientSession() as session:
+ arthur.instance = KingArthur(
+ guild_id=CONFIG.guild_id,
+ http_session=session,
+ command_prefix=commands.when_mentioned_or(*CONFIG.prefixes),
+ allowed_roles=(CONFIG.devops_role,),
+ case_insensitive=True,
+ intents=intents,
+ )
+ async with arthur.instance as bot:
+ await bot.start(CONFIG.token)
-if __name__ == "__main__":
- start()
+with arthur.logger.catch():
+ asyncio.run(main())
diff --git a/arthur/apis/cloudflare/zones.py b/arthur/apis/cloudflare/zones.py
index 81a54f1..7d407a0 100644
--- a/arthur/apis/cloudflare/zones.py
+++ b/arthur/apis/cloudflare/zones.py
@@ -8,30 +8,34 @@ from arthur.config import CONFIG
AUTH_HEADER = {"Authorization": f"Bearer {CONFIG.cloudflare_token}"}
-async def list_zones(zone_name: Optional[str] = None) -> dict[str, str]:
+async def list_zones(
+ session: aiohttp.ClientSession,
+ zone_name: Optional[str] = None,
+) -> dict[str, str]:
"""List all Cloudflare zones."""
endpoint = "https://api.cloudflare.com/client/v4/zones"
if zone_name is not None:
endpoint += f"?name={zone_name}"
- async with aiohttp.ClientSession() as session:
- async with session.get(endpoint, headers=AUTH_HEADER) as response:
- info = await response.json()
+ async with session.get(endpoint, headers=AUTH_HEADER) as response:
+ info = await response.json()
zones = info["result"]
return {zone["name"]: zone["id"] for zone in zones}
-async def purge_zone(zone_identifier: str) -> dict:
+async def purge_zone(
+ session: aiohttp.ClientSession,
+ zone_identifier: str,
+) -> dict:
"""Purge the cache for a Cloudflare zone."""
endpoint = f"https://api.cloudflare.com/client/v4/zones/{zone_identifier}/purge_cache"
request_body = {"purge_everything": True}
- async with aiohttp.ClientSession() as session:
- async with session.post(endpoint, headers=AUTH_HEADER, json=request_body) as response:
- info = await response.json()
+ async with session.post(endpoint, headers=AUTH_HEADER, json=request_body) as response:
+ info = await response.json()
return {"success": info["success"], "errors": info["errors"]}
diff --git a/arthur/bot.py b/arthur/bot.py
index 2bbb69c..2220db6 100644
--- a/arthur/bot.py
+++ b/arthur/bot.py
@@ -2,27 +2,20 @@
from pathlib import Path
from typing import Any, Union
+from botcore import BotBase
+from botcore.utils import scheduling
from discord import Interaction, Member, User
from discord.ext import commands
-from discord.ext.commands import Bot
from kubernetes_asyncio import config
-from arthur import logger
+from arthur import exts, logger
from arthur.config import CONFIG
-from arthur.extensions import find_extensions
-class KingArthur(Bot):
+class KingArthur(BotBase):
"""Base bot class for King Arthur."""
def __init__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None:
- config = {
- "command_prefix": commands.when_mentioned_or(*CONFIG.prefixes),
- "case_insensitive": True,
- }
-
- kwargs.update(config)
-
super().__init__(*args, **kwargs)
self.add_check(self._is_devops)
@@ -31,7 +24,10 @@ class KingArthur(Bot):
def _is_devops(ctx: Union[commands.Context, Interaction]) -> bool:
"""Check all commands are executed by authorised personnel."""
if isinstance(ctx, Interaction):
- return CONFIG.devops_role in [r.id for r in ctx.author.roles]
+ if isinstance(ctx.user, Member):
+ return CONFIG.devops_role in [r.id for r in ctx.user.roles]
+ else:
+ return False
if ctx.command.name == "ed":
return True
@@ -41,37 +37,23 @@ class KingArthur(Bot):
return CONFIG.devops_role in [r.id for r in ctx.author.roles]
- async def on_ready(self) -> None:
- """Initialise bot once connected and authorised with Discord."""
+ async def setup_hook(self) -> None:
+ """Async initialisation method for discord.py."""
+ await super().setup_hook()
+
# Authenticate with Kubernetes
if (Path.home() / ".kube/config").exists():
await config.load_kube_config()
else:
config.load_incluster_config()
-
logger.info(f"Logged in <red>{self.user}</>")
- # Start extension loading
-
- for path, extension in find_extensions():
- logger.info(
- f"Loading extension <magenta>{path.stem}</> " f"from <magenta>{path.parent}</>"
- )
-
- try:
- self.load_extension(extension)
- except: # noqa: E722
- logger.exception(
- f"Failed to load extension <magenta>{path.stem}</> "
- f"from <magenta>{path.parent}</>",
- )
- else:
- logger.info(
- f"Loaded extension <magenta>{path.stem}</> " f"from <magenta>{path.parent}</>"
- )
+ # This is not awaited to avoid a deadlock with any cogs that have
+ # wait_until_guild_available in their cog_load method.
+ scheduling.create_task(self.load_extensions(exts))
logger.info("Loading <red>jishaku</red>")
- self.load_extension("jishaku")
+ await self.load_extension("jishaku")
logger.info("Loaded <red>jishaku</red>")
async def is_owner(self, user: Union[User, Member]) -> bool:
diff --git a/arthur/config.py b/arthur/config.py
index a17980f..1fcf11d 100644
--- a/arthur/config.py
+++ b/arthur/config.py
@@ -17,6 +17,9 @@ class Config(BaseSettings):
# Token for authorising with the Cloudflare API
cloudflare_token: str
+ # Guild id
+ guild_id: int = 267624335836053506
+
class Config: # noqa: D106
env_file = ".env"
env_prefix = "KING_ARTHUR_"
diff --git a/arthur/extensions.py b/arthur/extensions.py
deleted file mode 100644
index 2170d68..0000000
--- a/arthur/extensions.py
+++ /dev/null
@@ -1,15 +0,0 @@
-"""Utilities for working with extensions."""
-from pathlib import Path
-from typing import Generator
-
-
-def find_extensions() -> Generator[tuple[Path, str], None, None]:
- """Search the exts directory to find cogs to load."""
- for path in Path("arthur/exts").rglob("**/*.py"):
- # Convert a path like "arthur/exts/foo/bar.py" to "arthur.exts.foo.bar"
- yield path, path_to_module(path)
-
-
-def path_to_module(path: Path) -> str:
- """Convert a path like "arthur/exts/foo/bar.py" to "arthur.exts.foo.bar"."""
- return str(path.parent.as_posix()).replace("/", ".") + f".{path.stem}"
diff --git a/arthur/exts/__init__.py b/arthur/exts/__init__.py
new file mode 100644
index 0000000..5428188
--- /dev/null
+++ b/arthur/exts/__init__.py
@@ -0,0 +1 @@
+"""Package of all extensions to load on bot startup."""
diff --git a/arthur/exts/cloudflare/__init__.py b/arthur/exts/cloudflare/__init__.py
new file mode 100644
index 0000000..f2fbac3
--- /dev/null
+++ b/arthur/exts/cloudflare/__init__.py
@@ -0,0 +1 @@
+"""Extensions related to Cloudflare."""
diff --git a/arthur/exts/cloudflare/zones.py b/arthur/exts/cloudflare/zones.py
index 7f9ca35..5840d36 100644
--- a/arthur/exts/cloudflare/zones.py
+++ b/arthur/exts/cloudflare/zones.py
@@ -1,4 +1,5 @@
"""The zones cog helps with managing Cloudflare zones."""
+import aiohttp
import discord
from discord.ext import commands
@@ -11,10 +12,11 @@ from arthur.utils import generate_error_message
class ZonesView(discord.ui.View):
"""This view allows users to select and purge the zones specified."""
- def __init__(self, domains: dict[str, str]) -> None:
+ def __init__(self, domains: dict[str, str], session: aiohttp.ClientSession) -> None:
super().__init__()
self.domains = domains
+ self.session = session
for domain, zone_id in self.domains.items():
self.children[0].add_option(label=domain, value=domain, description=zone_id, emoji="🌐")
@@ -31,13 +33,13 @@ class ZonesView(discord.ui.View):
placeholder="Select a zone to purge...",
)
async def select_zones(
- self, dropdown: discord.ui.Select, interaction: discord.Interaction
+ self, interaction: discord.Interaction, dropdown: discord.ui.Select
) -> None:
"""Drop down menu contains the list of zones."""
zone_name = dropdown.values[0]
required_id = self.domains[zone_name]
- purge_attempt_response = await zones.purge_zone(required_id)
+ purge_attempt_response = await zones.purge_zone(self.session, required_id)
if purge_attempt_response["success"]:
message = ":white_check_mark:"
message += " **Cache cleared!** "
@@ -69,12 +71,12 @@ class Zones(commands.Cog):
@zones.command(name="purge")
async def purge(self, ctx: commands.Context) -> None:
"""Command to clear the Cloudflare cache of the specified zone."""
- cf_zones = await zones.list_zones()
+ cf_zones = await zones.list_zones(self.bot.http_session)
- view = ZonesView(cf_zones)
+ view = ZonesView(cf_zones, self.bot.http_session)
await ctx.send(":cloud: Pick which zone(s) that should have their cache purged", view=view)
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add the extension to the bot."""
- bot.add_cog(Zones(bot))
+ await bot.add_cog(Zones(bot))
diff --git a/arthur/exts/error_handler/__init__.py b/arthur/exts/error_handler/__init__.py
new file mode 100644
index 0000000..af8bc5e
--- /dev/null
+++ b/arthur/exts/error_handler/__init__.py
@@ -0,0 +1 @@
+"""Error handling extensions."""
diff --git a/arthur/exts/error_handler/error_handler.py b/arthur/exts/error_handler/error_handler.py
index 673fbc7..14a8ed2 100644
--- a/arthur/exts/error_handler/error_handler.py
+++ b/arthur/exts/error_handler/error_handler.py
@@ -55,6 +55,6 @@ class ErrorHandler(Cog):
)
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add cog to bot."""
- bot.add_cog(ErrorHandler(bot))
+ await bot.add_cog(ErrorHandler(bot))
diff --git a/arthur/exts/fun/__init__.py b/arthur/exts/fun/__init__.py
new file mode 100644
index 0000000..9aada5d
--- /dev/null
+++ b/arthur/exts/fun/__init__.py
@@ -0,0 +1 @@
+"""Extensions made just for fun."""
diff --git a/arthur/exts/fun/ed.py b/arthur/exts/fun/ed.py
index 24e0025..33f5ed4 100644
--- a/arthur/exts/fun/ed.py
+++ b/arthur/exts/fun/ed.py
@@ -11,12 +11,12 @@ class Ed(Cog):
def __init__(self, bot: KingArthur) -> None:
self.bot = bot
- @command(name="ed", usage="[-GVhs] [-p string] [file]")
- async def ed(self, ctx: Context, *, _args: str) -> None:
+ @command(name="ed", help="[-GVhs] [-p string] [file]")
+ async def ed(self, ctx: Context) -> None:
"""Ed is the standard text editor."""
await ctx.send("?")
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add cog to bot."""
- bot.add_cog(Ed(bot))
+ await bot.add_cog(Ed(bot))
diff --git a/arthur/exts/kubernetes/__init__.py b/arthur/exts/kubernetes/__init__.py
new file mode 100644
index 0000000..2825b98
--- /dev/null
+++ b/arthur/exts/kubernetes/__init__.py
@@ -0,0 +1 @@
+"""Extensions relates to Kubernetes."""
diff --git a/arthur/exts/kubernetes/certificates.py b/arthur/exts/kubernetes/certificates.py
index 096b936..a031076 100644
--- a/arthur/exts/kubernetes/certificates.py
+++ b/arthur/exts/kubernetes/certificates.py
@@ -52,6 +52,6 @@ class Certificates(commands.Cog):
await ctx.send(return_message.format(namespace, table))
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add the extension to the bot."""
- bot.add_cog(Certificates(bot))
+ await bot.add_cog(Certificates(bot))
diff --git a/arthur/exts/kubernetes/deployments.py b/arthur/exts/kubernetes/deployments.py
index ec4f846..7a69baf 100644
--- a/arthur/exts/kubernetes/deployments.py
+++ b/arthur/exts/kubernetes/deployments.py
@@ -36,7 +36,7 @@ class ConfirmDeployment(ui.View):
return False
@ui.button(label="Confirm", style=ButtonStyle.green, row=0)
- async def confirm(self, _button: ui.Button, interaction: Interaction) -> None:
+ async def confirm(self, interaction: Interaction, _button: ui.Button) -> None:
"""Redeploy the specified service."""
try:
await deployments.restart_deployment(self.deployment, self.namespace)
@@ -66,7 +66,7 @@ class ConfirmDeployment(ui.View):
self.stop()
@ui.button(label="Cancel", style=ButtonStyle.grey, row=0)
- async def cancel(self, _button: ui.Button, interaction: Interaction) -> None:
+ async def cancel(self, interaction: Interaction, _button: ui.Button) -> None:
"""Logic for if the deployment is not approved."""
await interaction.message.edit(
content=":x: Redeployment aborted",
@@ -173,6 +173,6 @@ class Deployments(commands.Cog):
)
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add the extension to the bot."""
- bot.add_cog(Deployments(bot))
+ await bot.add_cog(Deployments(bot))
diff --git a/arthur/exts/kubernetes/jobs.py b/arthur/exts/kubernetes/jobs.py
index 6f7a1fe..6259f65 100644
--- a/arthur/exts/kubernetes/jobs.py
+++ b/arthur/exts/kubernetes/jobs.py
@@ -37,7 +37,7 @@ class CronJobView(discord.ui.View):
placeholder="Select a CronJob to trigger...",
)
async def select_job(
- self, dropdown: discord.ui.Select, interaction: discord.Interaction
+ self, interaction: discord.Interaction, dropdown: discord.ui.Select
) -> None:
"""Drop down menu contains the list of cronjobsb."""
cronjob_namespace, cronjob_name = dropdown.values[0].split("/")
@@ -76,6 +76,6 @@ class Jobs(commands.Cog):
await ctx.send(":tools: Pick a CronJob to trigger", view=view)
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add the extension to the bot."""
- bot.add_cog(Jobs(bot))
+ await bot.add_cog(Jobs(bot))
diff --git a/arthur/exts/kubernetes/nodes.py b/arthur/exts/kubernetes/nodes.py
index 1c58776..f42d3ed 100644
--- a/arthur/exts/kubernetes/nodes.py
+++ b/arthur/exts/kubernetes/nodes.py
@@ -102,6 +102,6 @@ class Nodes(commands.Cog):
)
-def setup(bot: KingArthur) -> None:
+async def setup(bot: KingArthur) -> None:
"""Add the extension to the bot."""
- bot.add_cog(Nodes(bot))
+ await bot.add_cog(Nodes(bot))
generated by cgit v1.2.3 (git 2.39.1) at 2025-11-07 08:57:51 +0000