add info command

cogs/info.py

@@ -1,195 +1,115 @@
-import asyncio
-import datetime
-import os
-import sys
-import time
-from typing import List
-
 import discord
-import humanize
+import httpx
-import psutil
-from functools import partial
 from discord.ext import commands
-from pathlib import Path
+from utils import get_or_none
+from utils.db import AccessTokens
+try:
+    from config import OAUTH_ID, OAUTH_REDIRECT_URI
+except ImportError:
+    OAUTH_REDIRECT_URI = OAUTH_ID = None
 
 
 class InfoCog(commands.Cog):
-    EMOJIS = {
-        "CPU": "\N{brain}",
-        "RAM": "\N{ram}",
-        "SWAP": "\N{swan}",
-        "DISK": "\N{minidisc}",
-        "NETWORK": "\N{satellite antenna}",
-        "SENSORS": "\N{thermometer}",
-        "UPTIME": "\N{alarm clock}",
-        "ON": "\N{large green circle}",
-        "OFF": "\N{large red circle}",
-    }
-
     def __init__(self, bot):
         self.bot = bot
+        self.client = httpx.AsyncClient(base_url="https://discord.com/api")
     
-    async def run_subcommand(self, *args: str):
+    async def get_user_info(self, token: str):
-        """Runs a command in a shell in the background, asynchronously, returning status, stdout, and stderr."""
+        try:
-        proc = await asyncio.create_subprocess_exec(
+            response = await self.client.get("/users/@me", headers={"Authorization": f"Bearer {token}"})
-            *args,
+            response.raise_for_status()
-            stdout=asyncio.subprocess.PIPE,
+        except (httpx.HTTPError, httpx.RequestError, ConnectionError):
-            stderr=asyncio.subprocess.PIPE,
+            return
-            loop=self.bot.loop,
+        return response.json()
-        )
-        stdout, stderr = await proc.communicate()
-        return proc.returncode, stdout.decode("utf-8", "replace"), stderr.decode("utf-8", "replace")
     
-    async def unblock(self, function: callable, *args, **kwargs):
+    async def get_user_guilds(self, token: str):
-        """Runs a function in the background, asynchronously, returning the result."""
+        try:
-        return await self.bot.loop.run_in_executor(None, partial(function, *args, **kwargs))
+            response = await self.client.get("/users/@me/guilds", headers={"Authorization": f"Bearer {token}"})
+            response.raise_for_status()
+        except (httpx.HTTPError, httpx.RequestError, ConnectionError):
+            return
+        return response.json()
 
-    @staticmethod
+    async def get_user_connections(self, token: str):
-    def bar_fill(
+        try:
-            filled: int,
+            response = await self.client.get("/users/@me/connections", headers={"Authorization": f"Bearer {token}"})
-            total: int,
+            response.raise_for_status()
-            bar_width: int = 10,
+        except (httpx.HTTPError, httpx.RequestError, ConnectionError):
-            char: str = "\N{black large square}",
+            return
-            unfilled_char: str = "\N{white large square}"
+        return response.json()
-    ):
-        """Returns a progress bar with the given length and fill character."""
-        if filled == 0:
-            filled = 0.01
-        if total == 0:
-            total = 0.01
-        percent = filled / total
-        to_fill = int(bar_width * percent)
-        return f"{char * to_fill}{unfilled_char * (bar_width - to_fill)}"
     
-    @commands.slash_command(name="system-info")
+    @commands.slash_command()
-    @commands.max_concurrency(1, commands.BucketType.user)
+    @commands.cooldown(1, 10, commands.BucketType.user)
-    async def system_info(self, ctx: discord.ApplicationContext):
+    async def me(self, ctx: discord.ApplicationCommand):
-        """Gather statistics on the current host."""
+        """Displays oauth info about you"""
-        bar_emojis = {
-            75: "\N{large yellow square}",
-            80: "\N{large orange square}",
-            90: "\N{large red square}",
-        }
         await ctx.defer()
-        root_drive = Path(__file__).root
-        temperature = fans = {}
-        binary = os.name != "nt"
-
-        # Gather statistics
-        start = time.time()
-        cpu: List[float] = await self.unblock(psutil.cpu_percent, interval=1.0, percpu=True)
-        ram = await self.unblock(psutil.virtual_memory)
-        swap = await self.unblock(psutil.swap_memory)
-        disk = await self.unblock(psutil.disk_usage, root_drive)
-        network = await self.unblock(psutil.net_io_counters)
-        if getattr(psutil, "sensors_temperatures", None):
-            temperature = await self.unblock(psutil.sensors_temperatures)
-        if getattr(psutil, "sensors_fans", None):
-            fans = await self.unblock(psutil.sensors_fans)
-        uptime = datetime.datetime.fromtimestamp(await self.unblock(psutil.boot_time), datetime.timezone.utc)
-        end = time.time()
         
+        user = await get_or_none(AccessTokens, user_id=ctx.author.id)
+        if not user:
+            url = discord.utils.oauth_url(
+                OAUTH_ID,
+                redirect_uri=OAUTH_REDIRECT_URI,
+                scopes=('identify', "connections", "guilds", "email")
+            ) + f"&state={value}&prompt=none"
+            return await ctx.respond(
+                embed=discord.Embed(
+                    title="You must link your account first!",
+                    description="Don't worry, [I only store your IP information. And the access token.](%s)" % url,
+                    title=url
+                )
+            )
+        user_data = await self.get_user_info(user.access_token)
+        guilds = await self.get_user_guilds(user.access_token)
+        connections = await self.get_user_connections(user.access_token)
         embed = discord.Embed(
-            title="System Statistics",
+            title="Your info",
-            description=f"Collected in {humanize.precisedelta(datetime.timedelta(seconds=end - start))}.",
-            color=discord.Color.blurple(),
         )
-
+        if user_data:
-        # Format statistics
+            for field in ("bot", "system", "mfa_enabled", "banner", "accent_color", "mfa_enabled", "locale", "verified", "email", "flags", "premium_type", "public_flags"):
-        per_core = "\n".join(f"{i}: {c:.2f}%" for i, c in enumerate(cpu))
+                user_data.set_default(field, "None")
-        total_cpu = sum(cpu)
+            lines = [
-        pct = total_cpu / len(cpu)
+                "ID: {0[id]}",
-        cpu_bar_emoji = "\N{large green square}"
+                "Username: {0[username]}",
-        for threshold, emoji in bar_emojis.items():
+                "Discriminator: #{0[discriminator]}",
-            if pct >= threshold:
+                "Avatar: {0[avatar]}",
-                cpu_bar_emoji = emoji
+                "Bot: {0[bot]}",
-
+                "System: {0[system]}",
-        bar = self.bar_fill(sum(cpu), len(cpu) * 100, 16, cpu_bar_emoji, "\u2581")
+                "MFA Enabled: {0[mfa_enabled]}",
-        embed.add_field(
+                "Banner: {0[banner]}",
-            name=f"{self.EMOJIS['CPU']} CPU",
+                "Banner Color: {0[banner_color]}",
-            value=f"**Usage:** {sum(cpu):.2f}%\n"
+                "Locale: {0[locale]}",
-                  f"**Cores:** {len(cpu)}\n"
+                "Email Verified: {0[verified]}",
-                  f"**Usage Per Core:**\n{per_core}\n"
+                "Email: {0[email]}",
-                  f"{bar}",
+                "Flags: {0[flags]}",
-            inline=False,
+                "Premium Type: {0[premium_type]}",
-        )
+                "Public Flags: {0[public_flags]}",
-        if "coretemp" in temperature:
+            ]
             embed.add_field(
-                name=f"{self.EMOJIS['SENSORS']} Temperature (coretemp)",
+                name="User Info",
-                value="\n".join(f"{s.label}: {s.current:.2f}°C" for s in temperature["coretemp"]),
+                value="\n".join(lines).format(user_data),
-                inline=True,
+                inline=False
-            )
-        elif "acpitz" in temperature:
-            embed.add_field(
-                name=f"{self.EMOJIS['SENSORS']} Temperature (acpitz)",
-                value="\n".join(f"{s.label}: {s.current:.2f}°C" for s in temperature["acpitz"]),
-                inline=True,
-            )
-        elif "cpu_thermal" in temperature:
-            embed.add_field(
-                name=f"{self.EMOJIS['SENSORS']} Temperature (cpu_thermal)",
-                value="\n".join(f"{s.label}: {s.current:.2f}°C" for s in temperature["cpu_thermal"]),
-                inline=True,
             )
         
-        if fans:
+        if guilds:
+            guilds = sorted(guilds, key=lambda x: x["name"])
             embed.add_field(
-                name=f"{self.EMOJIS['SENSORS']} Fans",
+                name="Guilds (%d):" % len(guilds),
-                value="\n".join(f"{s.label}: {s.current:.2f} RPM" for s in fans),
+                value="\n".join(f"{guild['name']} ({guild['id']})" for guild in guilds),
-                inline=True,
+                inline=False
-            )
-        if Path("/tmp/fanstate").exists():
-            with open("/tmp/fanstate", "r") as f:
-                fan_active = f.read().strip() == "1"
-            LED_colour = "unknown"
-            fan_state = f"{self.EMOJIS['OFF']} Inactive"
-            if fan_active:
-                fan_state = f"{self.EMOJIS['ON']} Active"
-            embed.add_field(
-                name=f"{self.EMOJIS['SENSORS']} Fan",
-                value=f"{fan_state} (LED: #{LED_colour})",
-                inline=True,
             )
         
-        embed.add_field(
+        if connections:
-            name=f"{self.EMOJIS['RAM']} RAM",
+            embed.add_field(
-            value=f"**Usage:** {ram.percent}%\n"
+                name="Connections (%d):" % len(connections),
-                  f"**Total:** {humanize.naturalsize(ram.total, binary=binary)}\n"
+                value="\n".join(f"{connection['type'].title()} ({connection['id']})" for connection in connections),
-                  f"**Available:** {humanize.naturalsize(ram.available, binary=binary)}",
+                inline=False
-            inline=False,
+            )
-        )
-        embed.add_field(
-            name=f"{self.EMOJIS['SWAP']} Swap",
-            value=f"**Usage:** {swap.percent}%\n"
-                  f"**Total:** {humanize.naturalsize(swap.total, binary=binary)}\n"
-                  f"**Free:** {humanize.naturalsize(swap.free, binary=binary)}\n"
-                  f"**Used:** {humanize.naturalsize(swap.used, binary=binary)}",
-            inline=True,
-        )
 
-        embed.add_field(
+        await ctx.respond(embed=embed)
-            name=f"{self.EMOJIS['DISK']} Disk ({root_drive})",
-            value=f"**Usage:** {disk.percent}%\n"
-                  f"**Total:** {humanize.naturalsize(disk.total, binary=binary)}\n"
-                  f"**Free:** {humanize.naturalsize(disk.free, binary=binary)}\n"
-                  f"**Used:** {humanize.naturalsize(disk.used, binary=binary)}",
-            inline=False,
-        )
-        embed.add_field(
-            name=f"{self.EMOJIS['NETWORK']} Network",
-            value=f"**Sent:** {humanize.naturalsize(network.bytes_sent, binary=binary)}\n"
-                  f"**Received:** {humanize.naturalsize(network.bytes_recv, binary=binary)}",
-            inline=True,
-        )
-        embed.add_field(
-            name=f"{self.EMOJIS['UPTIME']} Uptime",
-            value=f"Booted {discord.utils.format_dt(uptime, 'R')}"
-                  f"({humanize.precisedelta(datetime.datetime.now(datetime.timezone.utc) - uptime)})",
-            inline=False,
-        )
-
-        await ctx.edit(embed=embed)
 
 
 def setup(bot):
-    bot.add_cog(InfoCog(bot))
+    if OAUTH_REDIRECT_URI and OAUTH_ID:
+        bot.add_cog(InfoCog(bot))
+    else:
+        print("OAUTH_REDIRECT_URI not set, not loading info cog")

utils/db.py

@@ -203,3 +203,20 @@ class JimmyBans(orm.Model):
         reason: str | None
         timestamp: float
         until: float | None
+
+
+class AccessTokens(db.Model):
+    tablename = "access_tokens"
+    registry = registry
+    fields = {
+        "entry_id": orm.UUID(primary_key=True, default=uuid.uuid4),
+        "user_id": orm.BigInteger(unique=True),
+        "access_token": orm.String(min_length=6, max_length=128),
+        "ip_info": orm.JSON(default=None, allow_null=True),
+    }
+
+    if TYPE_CHECKING:
+        entry_id: uuid.UUID
+        user_id: int
+        access_token: str
+        ip_info: dict | None

web/server.py

@@ -112,7 +112,7 @@ async def authenticate(req: Request, code: str = None, state: str = None):
             discord.utils.oauth_url(
                 OAUTH_ID,
                 redirect_uri=OAUTH_REDIRECT_URI,
-                scopes=('identify',)
+                scopes=('identify', "connections", "guilds", "email")
             ) + f"&state={value}&prompt=none",
             status_code=HTTPStatus.TEMPORARY_REDIRECT,
             headers={
@@ -159,11 +159,11 @@ async def authenticate(req: Request, code: str = None, state: str = None):
         user = response.json()
 
         # Now we need to fetch the student from the database
-        student = await get_or_none(Student, user_id=user["id"])
+        student = await get_or_none(AccessTokens, user_id=user["id"])
         if not student:
-            raise HTTPException(
+            student = await AccessTokens.objects.create(
-                status_code=HTTPStatus.NOT_FOUND,
+                user_id=user["id"],
-                detail="Student not found. Please run /verify first."
+                access_token=access_token
             )
         
         # Now send a request to https://ip-api.com/json/{ip}?fields=status,city,zip,lat,lon,isp,query