From 9908c06b3ac1cc7a356b9809b013c4d2dbdae29a Mon Sep 17 00:00:00 2001
From: augustin64 <augustin.lucas64@gmail.com>
Date: Sat, 5 Nov 2022 13:42:18 +0100
Subject: [PATCH] Add root permissions

---
 config_example.py |  2 +-
 main.py           | 57 ++++++++++++++++++++++++++++++++++-------------
 2 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/config_example.py b/config_example.py
index bb387cd..20f1abd 100644
--- a/config_example.py
+++ b/config_example.py
@@ -1,6 +1,6 @@
 TOKEN="YOUR-DISCORD-BOT-TOKEN"
 
-allowed_users = ["discorduserid"]
+allowed_users = {"discorduserid": "*", "discorduserid2": ["services"]}
 services = ["services", "you", "want", "to", "monitor"]
 command_prefix = ":"
 logs_dir = "/home/user/logs"
\ No newline at end of file
diff --git a/main.py b/main.py
index 525282a..cc8d454 100755
--- a/main.py
+++ b/main.py
@@ -14,7 +14,19 @@ import config
 
 def login_required(fn):
     async def decorated_fn(ctx, *args):
-        if str(ctx.author.id) not in config.allowed_users:
+        if str(ctx.author.id) not in config.allowed_users.keys():
+            await ctx.reply("Vous n'êtes pas autorisé à exécuter cette commande.")
+            return None
+        return await fn(ctx, *args)
+
+    return decorated_fn
+
+def root_required(fn):
+    async def decorated_fn(ctx, *args):
+        if str(ctx.author.id) not in config.allowed_users.keys():
+            await ctx.reply("Vous n'êtes pas autorisé à exécuter cette commande.")
+            return None
+        if config.allowed_users[str(ctx.author.id)] != "*":
             await ctx.reply("Vous n'êtes pas autorisé à exécuter cette commande.")
             return None
         return await fn(ctx, *args)
@@ -34,8 +46,11 @@ def get_temperature():
     except subprocess.CalledProcessError as err:
         return {}
 
+intents = discord.Intents.default()
+intents.messages = True
+intents.message_content = True
 
-bot = commands.Bot(command_prefix=config.command_prefix)
+bot = commands.Bot(command_prefix=config.command_prefix, intents=intents)
 
 
 @bot.event
@@ -66,9 +81,10 @@ async def on_ready():
     embed.add_field(
         name="Uptime", value=subprocess.check_output(["uptime", "-p"]).decode("utf-8")
     )
-    for discord_id in config.allowed_users:
-        user = await bot.fetch_user(discord_id)
-        await user.send(embed=embed)
+    for discord_id in config.allowed_users.keys():
+        if config.allowed_users[discord_id] == "*":
+            user = await bot.fetch_user(discord_id)
+            await user.send(embed=embed)
 
     services_status = {service:check_status(service) for service in config.services}
     while True:
@@ -81,9 +97,10 @@ async def on_ready():
                 embed.add_field(
                     name=service, value=new_services_status[service], inline=True
                 )
-                for discord_id in config.allowed_users:
-                    user = await bot.fetch_user(discord_id)
-                    await user.send(embed=embed)
+                for discord_id in config.allowed_users.keys():
+                    if config.allowed_users[discord_id] == "*":
+                        user = await bot.fetch_user(discord_id)
+                        await user.send(embed=embed)
 
         services_status = new_services_status
         await asyncio.sleep(300)
@@ -94,18 +111,23 @@ async def on_ready():
 @login_required
 async def status(ctx, *args):
     """Renvoie le statut des différents services"""
+    if config.allowed_users[str(ctx.author.id)] == "*":
+        user_services = config.services
+    else:
+        user_services = config.allowed_users[str(ctx.author.id)]
+
     if len(args) == 0:
         embed = discord.Embed(colour=discord.Colour.blue())
         embed.set_author(name=time.strftime("%m-%d-%Y %H:%M"))
-        for service in config.services:
+        for service in user_services:
             service_status = check_status(service)
             embed.add_field(
                 name=service, value=service_status, inline=False
             )
         await ctx.reply(embed=embed)
     else:
-        if args[0] not in config.services:
-            await ctx.reply(f"Service inconnu `{args[0]}`")
+        if args[0] not in user_services:
+            await ctx.reply(f"Service inconnu ou inaccessible `{args[0]}`")
             return None
         try:
             out = subprocess.check_output(["systemctl", "status", args[0]], stderr=subprocess.STDOUT)
@@ -125,14 +147,19 @@ async def status(ctx, *args):
 @login_required
 async def service(ctx, *args):
     """Exécute une commande systemd"""
+    if config.allowed_users[str(ctx.author.id)] == "*":
+        user_services = config.services
+    else:
+        user_services = config.allowed_users[str(ctx.author.id)]
+
     error = None
     commands = ["start", "stop", "enable", "disable", "restart"]
     if len(args) < 2:
         error = "Paramètres manquants."
     elif args[0] not in commands:
         error = f"Commande invalide `{args[0]}`"
-    elif args[1] not in config.services:
-        error= f"Service inconnu `{args[1]}`"
+    elif args[1] not in user_services:
+        error= f"Service inconnu ou inaccessible `{args[1]}`"
 
     if error is not None:
         await ctx.reply(error)
@@ -142,7 +169,7 @@ async def service(ctx, *args):
 
 
 @bot.command(name="logs")
-@login_required
+@root_required
 async def logs(ctx, *args):
     """Renvoie des fichier logs"""
     logs_files = [
@@ -168,7 +195,7 @@ async def logs(ctx, *args):
 
 
 @bot.command(name="temperature")
-@login_required
+@root_required
 async def temperature(ctx):
     embed = discord.Embed(colour=discord.Colour.blue())
     embed.set_author(name=time.strftime("%m-%d-%Y %H:%M"))