From f10ab7bb508e067d5e84a63d0c7880c75eaaf0c5 Mon Sep 17 00:00:00 2001
From: TSRBerry <20988865+TSRBerry@users.noreply.github.com>
Date: Thu, 30 Mar 2023 19:01:04 +0200
Subject: [PATCH] Add role persistence cog (#23)

---
 robocop_ng/__main__.py              |  1 +
 robocop_ng/cogs/role_persistence.py | 29 ++++++++++++++++++++++++
 robocop_ng/helpers/roles.py         | 34 +++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 robocop_ng/cogs/role_persistence.py
 create mode 100644 robocop_ng/helpers/roles.py

diff --git a/robocop_ng/__main__.py b/robocop_ng/__main__.py
index 799ed71..0ba16d4 100755
--- a/robocop_ng/__main__.py
+++ b/robocop_ng/__main__.py
@@ -46,6 +46,7 @@ wanted_jsons = [
     "data/robocronptab.json",
     "data/userlog.json",
     "data/invites.json",
+    "data/persistent_roles.json"
 ]
 
 intents = discord.Intents.all()
diff --git a/robocop_ng/cogs/role_persistence.py b/robocop_ng/cogs/role_persistence.py
new file mode 100644
index 0000000..cdb9e17
--- /dev/null
+++ b/robocop_ng/cogs/role_persistence.py
@@ -0,0 +1,29 @@
+from discord import RawMemberRemoveEvent, Member
+from discord.ext.commands import Cog
+
+from robocop_ng.helpers.roles import add_user_roles, get_user_roles
+
+
+class RolePersistence(Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+    @Cog.listener
+    async def on_raw_member_remove(self, payload: RawMemberRemoveEvent):
+        save_roles = []
+        for role in payload.user.roles:
+            if role.is_assignable() and \
+                    not role.is_default() and not role.is_premium_subscriber() and \
+                    not role.is_bot_managed() and not role.is_integration():
+                save_roles.append(role.id)
+
+        if len(save_roles) > 0:
+            add_user_roles(payload.user.id, save_roles)
+
+    @Cog.listener
+    async def on_member_join(self, member: Member):
+        user_roles = get_user_roles(member.id)
+        if len(user_roles) > 0:
+            user_roles = [member.guild.get_role(int(role)) for role in user_roles
+                          if member.guild.get_role(int(role)) is not None]
+            await member.add_roles(*user_roles, reason="Restoring old roles from `role_persistence`.")
diff --git a/robocop_ng/helpers/roles.py b/robocop_ng/helpers/roles.py
new file mode 100644
index 0000000..50830fe
--- /dev/null
+++ b/robocop_ng/helpers/roles.py
@@ -0,0 +1,34 @@
+import json
+import os.path
+
+PERSISTENT_ROLES_FILE = "data/persistent_roles.json"
+
+
+def get_persistent_roles() -> dict[str, list[str]]:
+    if os.path.isfile(PERSISTENT_ROLES_FILE):
+        with open(PERSISTENT_ROLES_FILE, "r") as f:
+            return json.load(f)
+    return {}
+
+
+def set_persistent_roles(contents: dict[str, list[str]]):
+    with open(PERSISTENT_ROLES_FILE, "w") as f:
+        json.dump(contents, f)
+
+
+def add_user_roles(uid: int, roles: list[int]):
+    uid = str(uid)
+    roles = [str(x) for x in roles]
+
+    persistent_roles = get_persistent_roles()
+    persistent_roles[uid] = roles
+    set_persistent_roles(persistent_roles)
+
+
+def get_user_roles(uid: int) -> list[str]:
+    uid = str(uid)
+    with open(PERSISTENT_ROLES_FILE, "r") as f:
+        roles = json.load(f)
+        if uid in roles:
+            return roles[uid]
+        return []