From 959d83f0b121f9ae3fa0000d7fb6519245246a66 Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Wed, 1 May 2024 00:46:03 +0100
Subject: [PATCH] Allow configuring IP servers

---
 config.example.toml |  6 ++++++
 src/cogs/net.py     | 36 +++++++++++++++++++++++++++++++-----
 src/conf.py         |  1 +
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/config.example.toml b/config.example.toml
index 45efb89..fc2360c 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -71,3 +71,9 @@ downloads_pdfs = true  # If False, PDF links will not be downloaded and uploaded
 enabled = true  # whether any transcoding at all should be done. If false, no links or attachments are ever probed.
 hevc = true  # Enables probing video links and converting any HEVC videos to H264+opus.
 on_demand = true  # Enables transcoding any video to h264+opus when 📼 (VHS) is reacted.
+
+[network]
+ip_servers = [
+    "https://ip.i-am.nexus",
+    "https://ip.shronk.tech"
+]  # A list of SHRoNK IP servers. See: https://github.com/SHRoNK-Corporation/shronk-ip-spec/blob/main/RFS0001.md
diff --git a/src/cogs/net.py b/src/cogs/net.py
index 21081f6..188f096 100644
--- a/src/cogs/net.py
+++ b/src/cogs/net.py
@@ -13,11 +13,13 @@ from discord.ext import commands
 from dns import asyncresolver
 from rich.console import Console
 from rich.tree import Tree
+from conf import CONFIG
 
 
 class NetworkCog(commands.Cog):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
+        self.config = CONFIG["network"]
 
     @commands.slash_command()
     async def ping(self, ctx: discord.ApplicationContext, target: str = None):
@@ -255,6 +257,21 @@ class NetworkCog(commands.Cog):
             await ctx.defer()
             await ctx.respond(file=discord.File(f))
 
+    async def _fetch_ip_response(
+            self,
+            server: str,
+            lookup: str,
+            client: httpx.AsyncClient
+    ) -> tuple[dict, float] | httpx.HTTPError | ConnectionError | json.JSONDecodeError:
+        try:
+            start = time.perf_counter()
+            response = await client.get(f"https://{server}/lookup?ip={lookup}")
+            end = time.perf_counter()
+        except (httpx.HTTPError, ConnectionError) as e:
+            return e
+        else:
+            return response.json(), round(end - start, 2)
+
     @commands.slash_command(name="ip")
     async def get_ip_address(self, ctx: discord.ApplicationContext, lookup: str = None):
         """Fetches IP info from SHRONK IP servers"""
@@ -263,14 +280,23 @@ class NetworkCog(commands.Cog):
             if not lookup:
                 response = await client.get("https://api.ipify.org")
                 lookup = response.text
-            servers = [
-                "ip.shronk.net",
-                "ip.i-am.nexus",
-                "ip.shronk.nicroxio.co.uk"
-            ]
+
+            servers = self.config.get(
+                "ip_servers",
+                [
+                    "ip.shronk.net",
+                    "ip.i-am.nexus",
+                    "ip.shronk.nicroxio.co.uk"
+                ]
+            )
             embed = discord.Embed(
                 title="IP lookup information for: %s" % lookup,
             )
+            if not servers:
+                embed.description = "No IP servers configured in config.toml."
+                embed.colour = discord.Color.red()
+                return await ctx.respond(embed=embed)
+
             for server in servers:
                 try:
                     start = time.perf_counter()
diff --git a/src/conf.py b/src/conf.py
index 04b0d7c..111b67a 100644
--- a/src/conf.py
+++ b/src/conf.py
@@ -39,6 +39,7 @@ CONFIG.setdefault("jimmy", {})
 CONFIG.setdefault("ollama", {})
 CONFIG.setdefault("screenshot", {})
 CONFIG.setdefault("responder", {})
+CONFIG.setdefault("network", {})
 CONFIG.setdefault("quote_a", {"channel": None})
 CONFIG.setdefault("redis", {"host": "redis", "port": 6379, "decode_responses": True})