From 3c63f26e79d2be92511cccae298305b38ec74793 Mon Sep 17 00:00:00 2001
From: zevaryx <zevaryx@gmail.com>
Date: Sun, 20 Mar 2022 14:33:17 -0600
Subject: [PATCH] Add Lock and Lockdown automated tasks

---
 jarvis_tasks/__init__.py       | 11 ++++++--
 jarvis_tasks/tasks/lock.py     | 46 ++++++++++++++++++++++++++++++++++
 jarvis_tasks/tasks/lockdown.py | 41 ++++++++++++++++++++++++++++++
 pyproject.toml                 |  2 +-
 4 files changed, 97 insertions(+), 3 deletions(-)
 create mode 100644 jarvis_tasks/tasks/lock.py
 create mode 100644 jarvis_tasks/tasks/lockdown.py

diff --git a/jarvis_tasks/__init__.py b/jarvis_tasks/__init__.py
index 35d1dbb..b0fbf99 100644
--- a/jarvis_tasks/__init__.py
+++ b/jarvis_tasks/__init__.py
@@ -7,7 +7,7 @@ from jarvis_core.db import connect
 from jarvis_core.log import get_logger
 
 from jarvis_tasks.config import TaskConfig
-from jarvis_tasks.tasks import ban, reminder, twitter, warning
+from jarvis_tasks.tasks import ban, lock, lockdown, reminder, twitter, warning
 
 logger = None
 
@@ -40,7 +40,14 @@ async def _start(config: Optional[str] = "config.yaml") -> None:
     # Start tasks
     try:
         logger.debug("Starting tasks")
-        functions = [ban.unban, reminder.remind, twitter.twitter, warning.unwarn]
+        functions = [
+            ban.unban,
+            lock.unlock,
+            lockdown.lift,
+            reminder.remind,
+            twitter.twitter,
+            warning.unwarn,
+        ]
         tasks = [loop.create_task(f(bot, logger)) for f in functions]
         for task in tasks:
             await task
diff --git a/jarvis_tasks/tasks/lock.py b/jarvis_tasks/tasks/lock.py
new file mode 100644
index 0000000..744811c
--- /dev/null
+++ b/jarvis_tasks/tasks/lock.py
@@ -0,0 +1,46 @@
+"""JARVIS Lock background task."""
+import asyncio
+from datetime import datetime, timedelta
+from logging import Logger
+
+from dis_snek import Snake
+from dis_snek.client.utils.misc_utils import get
+from dis_snek.models.discord.channel import GuildChannel
+from jarvis_core.db import q
+from jarvis_core.db.models import Lock
+
+from jarvis_tasks.util import runat
+
+
+async def _unlock(channel: GuildChannel, lock: Lock, logger: Logger) -> None:
+    logger.debug(f"Deactivating lock {lock.id}")
+    overwrite = get(channel.permission_overwrites, id=lock.guild)
+    if overwrite and lock.original_perms:
+        overwrite.allow = lock.original_perms.allow
+        overwrite.deny = lock.original_perms.deny
+        await channel.edit_permission(overwrite, reason="Automatic unlock")
+    elif overwrite and not lock.original_perms:
+        await channel.delete_permission(target=overwrite, reason="Automatic unlock")
+    else:
+        logger.debug("Permission neither exists not existed")
+
+
+async def unlock(bot: Snake, logger: Logger) -> None:
+    """
+    Unlock locked channels.
+
+    Args:
+        bot: Bot instance
+        logger: Global logger
+    """
+    while True:
+        max_ts = datetime.utcnow() + timedelta(seconds=55)
+        locks = Lock.find(q(active=True, created_at__lte=max_ts))
+        async for lock in locks:
+            guild = await bot.fetch_guild(lock.guild)
+            channel = await guild.fetch_channel(lock.channel)
+            coro = _unlock(channel, lock, logger)
+            when = lock.created_at + timedelta(minutes=lock.duration)
+            asyncio.create_task(runat(when, coro, logger))
+
+        await asyncio.sleep(delay=60)
diff --git a/jarvis_tasks/tasks/lockdown.py b/jarvis_tasks/tasks/lockdown.py
new file mode 100644
index 0000000..36e66b1
--- /dev/null
+++ b/jarvis_tasks/tasks/lockdown.py
@@ -0,0 +1,41 @@
+"""JARVIS Lockdown background task."""
+import asyncio
+from datetime import datetime, timedelta
+from logging import Logger
+
+from dis_snek import Snake
+from dis_snek.models.discord.enums import Permissions
+from dis_snek.models.discord.role import Role
+from jarvis_core.db import q
+from jarvis_core.db.models import Lockdown
+
+from jarvis_tasks.util import runat
+
+
+async def _lift(role: Role, lock: Lockdown, logger: Logger) -> None:
+    logger.debug(f"Lifting lockdown {lock.id}")
+    original_perms = Permissions(lock.original_perms)
+    await role.edit(permissions=original_perms)
+    lock.active = False
+    await lock.commit()
+
+
+async def lift(bot: Snake, logger: Logger) -> None:
+    """
+    Lift a lockdown.
+
+    Args:
+        bot: Bot instance
+        logger: Global logger
+    """
+    while True:
+        max_ts = datetime.utcnow() + timedelta(seconds=55)
+        locks = Lockdown.find(q(active=True, created_at__lte=max_ts))
+        async for lock in locks:
+            guild = await bot.fetch_guild(lock.guild)
+            role = await guild.fetch_role(guild.id)
+            coro = _lift(role, lock, logger)
+            when = lock.created_at + timedelta(minutes=lock.duration)
+            asyncio.create_task(runat(when, coro, logger))
+
+        await asyncio.sleep(delay=60)
diff --git a/pyproject.toml b/pyproject.toml
index f25dc7e..ddce357 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jarvis-tasks"
-version = "0.2.0"
+version = "0.3.0"
 description = ""
 authors = ["Your Name <you@example.com>"]