Added <prefix>Stream and <prefix>revokestream commands

3 files changed, 169 insertions, 0 deletions

diff --git a/bot/constants.py b/bot/constants.py
index 2126b2b37..744fbd512 100644
--- a/bot/constants.py
+++ b/bot/constants.py
@@ -466,6 +466,7 @@ class Roles(metaclass=YAMLGetter):
     unverified: int
     verified: int  # This is the Developers role on PyDis, here named verified for readability reasons.
     voice_verified: int
+    video: int
 
 
 class Guild(metaclass=YAMLGetter):
@@ -701,3 +702,30 @@ ERROR_REPLIES = [
     "Noooooo!!",
     "I can't believe you've done this",
 ]
+
+# TIME_FORMATS defines aliases and multipliers for time formats
+# key is a standard time unit name like second ,year, decade etc.
+# mul is a multiplier where duration of said time unit * multiplier = time in seconds
+# eg. 1 day = 1 * multiplier seconds, so mul = 86400
+TIME_FORMATS = {
+    "second": {
+        "aliases": ("s", "sec", "seconds", "secs"),
+        "mul": 1
+    },
+    "minute": {
+        "aliases": ("m", "min", "mins", "minutes"),
+        "mul": 60
+    },
+    "hour": {
+        "aliases": ("h", "hr", "hrs", "hours"),
+        "mul": 3600
+    },
+    "day": {
+        "aliases": ("d", "days"),
+        "mul": 86400
+    },
+    "year": {
+        "aliases": ("yr", "yrs", "years"),
+        "mul": 31536000
+    }
+}
diff --git a/bot/exts/moderation/stream.py b/bot/exts/moderation/stream.py
new file mode 100644
index 000000000..673a21b1b
--- /dev/null
+++ b/bot/exts/moderation/stream.py
@@ -0,0 +1,138 @@
+from discord.ext import commands, tasks
+import discord
+
+from bot.constants import Roles, STAFF_ROLES, Guild, TIME_FORMATS
+from bot import Bot
+import time
+from async_rediscache import RedisCache
+
+# Constant error messages
+NO_USER_SPECIFIED = "Please specify a user"
+TIME_FORMAT_NOT_VALID = "Please specify a valid time format ex. 10h or 1day"
+TIME_LESS_EQ_0 = "Duration can not be a 0 or lower"
+USER_ALREADY_ALLOWED_TO_STREAM = "This user can already stream"
+USER_ALREADY_NOT_ALLOWED_TO_STREAM = "This user already can't stream"
+
+
+# FORMATS holds a combined list of all allowed time units
+# made from TIME_FORMATS constant
+FORMATS = []
+for key, entry in TIME_FORMATS.items():
+    FORMATS.extend(entry["aliases"])
+    FORMATS.append(key)
+
+
+class Stream(commands.Cog):
+    """Stream class handles giving screen sharing permission with commands"""
+
+    # Data cache storing userid to unix_time relation
+    # user id is used to get member who's streaming permission need to be revoked after some time
+    # unix_time is a time when user's streaming permission needs tp be revoked in unix time notation
+    user_cache = RedisCache()
+
+    def __init__(self, bot: Bot):
+        self.bot = bot
+        self.remove_permissions.start()
+        self.guild_static = None
+
+    @staticmethod
+    def _link_from_alias(time_format) -> (dict, str):
+        """Get TIME_FORMATS key and entry by time format or any of its aliases"""
+        for format_key, val in TIME_FORMATS.items():
+            if format_key == time_format or time_format in val["aliases"]:
+                return TIME_FORMATS[format_key], format_key
+
+    def _parse_time_to_seconds(self, duration, time_format) -> int:
+        """Get time in seconds from duration and time format"""
+        return duration * self._link_from_alias(time_format)[0]["mul"]
+
+    @commands.command(aliases=("streaming", "share"))
+    @commands.has_any_role(*STAFF_ROLES)
+    async def stream(
+            self,
+            ctx: commands.Context,
+            user: discord.Member = None,
+            duration: int = 1,
+            time_format: str = "h",
+            *_
+    ):
+        """
+        stream handles <prefix>stream command
+        argument user - required user mention, any errors should be handled by upper level handler
+        duration - int must be higher than 0 - defaults to 1
+        time_format - str defining what time unit you want to use, must be any of FORMATS - defaults to h
+
+        Command give user permission to stream and takes it away after provided duration
+        """
+        # Check for required user argument
+        # if not provided send NO_USER_SPECIFIED message
+        if not user:
+            await ctx.send(NO_USER_SPECIFIED)
+            return
+
+        # Time can't be negative lol
+        if duration <= 0:
+            await ctx.send(TIME_LESS_EQ_0)
+            return
+
+        # Check if time_format argument is a valid time format
+        # eg. d, day etc are aliases for day time format
+        if time_format not in FORMATS:
+            await ctx.send(TIME_FORMAT_NOT_VALID)
+            return
+
+        # Check if user already has streaming permission
+        already_allowed = any(Roles.video == role.id for role in user.roles)
+        if already_allowed:
+            await ctx.send(USER_ALREADY_ALLOWED_TO_STREAM)
+            return
+
+        # Set user id - time in redis cache and add streaming permission role
+        await self.user_cache.set(user.id, time.time() + self._parse_time_to_seconds(duration, time_format))
+        await user.add_roles(discord.Object(Roles.video), reason="Temporary streaming access granted")
+        await ctx.send(f"{user.mention} can now stream for {duration} {self._link_from_alias(time_format)[1]}/s")
+
+    @tasks.loop(seconds=30)
+    async def remove_permissions(self):
+        """
+        background loop for removing streaming permission
+        """
+        all_entries = await self.user_cache.items()
+        for user_id, delete_time in all_entries:
+            if time.time() > delete_time:
+                member = self.guild_static.fetch_memebr(user_id)
+                if member:
+                    await member.remove_roles(discord.Object(Roles.video), reason="Temporary streaming access revoked")
+                    await self.user_cache.pop(user_id)
+
+    @remove_permissions.before_loop
+    async def await_ready(self):
+        """Wait for bot to be ready before starting remove_permissions loop
+        and get guild by id
+        """
+        await self.bot.wait_until_ready()
+        self.guild_static = self.bot.get_guild(Guild.id)
+
+    @commands.command(aliases=("unstream", ))
+    @commands.has_any_role(*STAFF_ROLES)
+    async def revokestream(
+            self,
+            ctx: commands.Context,
+            user: discord.Member = None
+    ):
+        """
+        stream handles <prefix>revokestream command
+        argument user - required user mention, any errors should be handled by upper level handler
+
+        command removes streaming permission from a user
+        """
+        not_allowed = not any(Roles.video == role.id for role in user.roles)
+        if not_allowed:
+            await user.remove_roles(discord.Object(Roles.video))
+        else:
+            await ctx.send(USER_ALREADY_NOT_ALLOWED_TO_STREAM)
+
+
+def setup(bot: Bot) -> None:
+    """Loads the Stream cog."""
+    bot.add_cog(Stream(bot))
diff --git a/config-default.yml b/config-default.yml
index 89493c4de..700406f4e 100644
--- a/config-default.yml
+++ b/config-default.yml
@@ -251,6 +251,9 @@ guild:
         jammers:        737249140966162473
         team_leaders:   737250302834638889
 
+        # Streaming
+        video: 764245844798079016
+
     moderation_roles:
         - *OWNERS_ROLE
         - *ADMINS_ROLE