add isolate, fixed nitro booster for punish

isolate/__init__.py

@@ -0,0 +1,7 @@
+from .isolate import Isolate
+
+
+async def setup(bot):
+    isolate = Isolate(bot)
+    await isolate.initialize()
+    bot.add_cog(isolate)

isolate/info.json

@@ -0,0 +1,20 @@
+{
+  "author": [
+    "Brandons209",
+    "calebj"
+  ],
+  "bot_version": [
+    3,
+    0,
+    0
+  ],
+  "description": "Provides a way to quickly remove a user for all chats and isolate them to a single channel. Supports setting times, automatic removal of isolate, modlog integration, custom role overrides, automatic server setup, and a lot more. This is different from punish as it always remove all a user's roles and by default makes the user not read all channels.",
+  "hidden": false,
+  "install_msg": "Thank you for using this cog! Please make sure to setup the cog using [p]isolateset",
+  "requirements": [],
+  "short": "Provides a way to quickly remove a user for all chats and isolate them to a single channel.",
+  "tags": [
+    "brandons209",
+    "isolate"
+  ]
+}

isolate/memoizer.py

@@ -0,0 +1,40 @@
+class Memoizer:
+    """
+    General purpose cache for function results. Appends positional args, overlays kwargs. Both must be hashable.
+    """
+
+    __slots__ = ["_cache", "_func", "_args", "_kwargs"]
+
+    def __init__(self, func, *args, **kwargs):
+        self._cache = {}
+        self._func = func
+        self._args = args
+        self._kwargs = kwargs
+
+    def clear(self):
+        "clears the internal arg -> result cache"
+        self._cache.clear()
+
+    def filter(self, iterable, *, skip_nulls=False):
+        "Calls the function on each item in the passed iterable. Only one positional arg at a time is supported."
+        gen = map(self, iterable)
+
+        if skip_nulls:
+            return filter(None, gen)
+        else:
+            return gen
+
+    def __call__(self, *args, **kwargs):
+        if kwargs:
+            key = (*args, tuple(kwargs.items()))
+            call_kwargs = self._kwargs.copy()
+            call_kwargs.update(kwargs)
+        else:
+            key = args
+            call_kwargs = self._kwargs
+
+        try:
+            return self._cache[key]
+        except KeyError:
+            ret = self._cache[key] = self._func(*self._args, *args, **call_kwargs)
+            return ret

isolate/utils.py

@@ -0,0 +1,221 @@
+import re
+import discord
+
+UNIT_TABLE = (
+    (("weeks", "wks", "w"), 60 * 60 * 24 * 7),
+    (("days", "dys", "d"), 60 * 60 * 24),
+    (("hours", "hrs", "h"), 60 * 60),
+    (("minutes", "mins", "m"), 60),
+    (("seconds", "secs", "s"), 1),
+)
+
+
+class BadTimeExpr(Exception):
+    pass
+
+
+def _find_unit(unit):
+    for names, length in UNIT_TABLE:
+        if any(n.startswith(unit) for n in names):
+            return names, length
+    raise BadTimeExpr("Invalid unit: %s" % unit)
+
+
+def parse_time(time):
+    time = time.lower()
+    if not time.isdigit():
+        time = re.split(r"\s*([\d.]+\s*[^\d\s,;]*)(?:[,;\s]|and)*", time)
+        time = sum(map(_timespec_sec, filter(None, time)))
+    return int(time)
+
+
+def _timespec_sec(expr):
+    atoms = re.split(r"([\d.]+)\s*([^\d\s]*)", expr)
+    atoms = list(filter(None, atoms))
+
+    if len(atoms) > 2:  # This shouldn't ever happen
+        raise BadTimeExpr("invalid expression: '%s'" % expr)
+    elif len(atoms) == 2:
+        names, length = _find_unit(atoms[1])
+        if atoms[0].count(".") > 1 or not atoms[0].replace(".", "").isdigit():
+            raise BadTimeExpr("Not a number: '%s'" % atoms[0])
+    else:
+        names, length = _find_unit("seconds")
+
+    try:
+        return float(atoms[0]) * length
+    except ValueError:
+        raise BadTimeExpr("invalid value: '%s'" % atoms[0])
+
+
+def generate_timespec(sec: int, short=False, micro=False) -> str:
+    timespec = []
+    sec = int(sec)
+    neg = sec < 0
+    sec = abs(sec)
+
+    for names, length in UNIT_TABLE:
+        n, sec = divmod(sec, length)
+
+        if n:
+            if micro:
+                s = "%d%s" % (n, names[2])
+            elif short:
+                s = "%d%s" % (n, names[1])
+            else:
+                s = "%d %s" % (n, names[0])
+
+            if n <= 1 and not (micro and names[2] == "s"):
+                s = s.rstrip("s")
+
+            timespec.append(s)
+
+    if len(timespec) > 1:
+        if micro:
+            spec = "".join(timespec)
+
+        segments = timespec[:-1], timespec[-1:]
+        spec = " and ".join(", ".join(x) for x in segments)
+    elif timespec:
+        spec = timespec[0]
+    else:
+        return "0"
+
+    if neg:
+        spec += " ago"
+
+    return spec
+
+
+def format_list(*items, join="and", delim=", "):
+    if len(items) > 1:
+        return (" %s " % join).join((delim.join(items[:-1]), items[-1]))
+    elif items:
+        return items[0]
+    else:
+        return ""
+
+
+def overwrite_to_dict(overwrite):
+    allow, deny = overwrite.pair()
+    return {"allow": allow.value, "deny": deny.value}
+
+
+def format_permissions(permissions, include_null=False):
+    entries = []
+
+    for perm, value in sorted(permissions, key=lambda t: t[0]):
+        if value is True:
+            symbol = "\N{WHITE HEAVY CHECK MARK}"
+        elif value is False:
+            symbol = "\N{NO ENTRY SIGN}"
+        elif include_null:
+            symbol = "\N{RADIO BUTTON}"
+        else:
+            continue
+
+        entries.append(symbol + " " + perm.replace("_", " ").title().replace("Tts", "TTS"))
+
+    if entries:
+        return "\n".join(entries)
+    else:
+        return "No permission entries."
+
+
+def getmname(mid, guild):
+    member = discord.utils.get(guild.members, id=int(mid))
+
+    if member:
+        return str(member)
+    else:
+        return "(absent user #%s)" % mid
+
+
+def role_from_string(guild, rolename, roles=None):
+    if rolename is None:
+        return None
+
+    if roles is None:
+        roles = guild.roles
+    else:
+        roles = [r for r in roles if isinstance(r, discord.Role)]
+
+    if type(rolename) == int:
+        role = discord.utils.get(roles, id=rolename)
+
+        if role:
+            return role
+
+    rolename = rolename.lower()
+    role = discord.utils.find(lambda r: r.name.lower() == rolename, roles)
+
+    return role
+
+
+def resolve_role_list(guild: discord.guild, roles: list) -> list:
+    gen = (role_from_string(guild, name) for name in roles)
+    return list(filter(None, gen))
+
+
+def permissions_for_roles(channel, *roles):
+    """
+    Calculates the effective permissions for a role or combination of roles.
+    Naturally, if no roles are given, the default role's permissions are used
+    """
+    default = channel.guild.default_role
+    base = discord.Permissions(default.permissions.value)
+
+    # Apply all role values
+    for role in roles:
+        base.value |= role.permissions.value
+
+    # guild-wide Administrator -> True for everything
+    # Bypass all channel-specific overrides
+    if base.administrator:
+        return discord.Permissions.all()
+
+    role_ids = set(map(lambda r: r.id, roles))
+    denies = 0
+    allows = 0
+
+    # Apply channel specific role permission overwrites
+    for target, overwrite in channel.overwrites.items():
+        # Handle default role first, if present
+        if overwrite.id == default.id:
+            base.handle_overwrite(allow=overwrite.allow, deny=overwrite.deny)
+
+        if isinstance(target, discord.Role) and target.id in role_ids:
+            denies |= overwrite.deny
+            allows |= overwrite.allow
+
+    base.handle_overwrite(allow=allows, deny=denies)
+
+    # default channels can always be read
+    if channel.is_default:
+        base.read_messages = True
+
+    # if you can't send a message in a channel then you can't have certain
+    # permissions as well
+    if not base.send_messages:
+        base.send_tts_messages = False
+        base.mention_everyone = False
+        base.embed_links = False
+        base.attach_files = False
+
+    # if you can't read a channel then you have no permissions there
+    if not base.read_messages:
+        denied = discord.Permissions.all_channel()
+        base.value &= ~denied.value
+
+    # text channels do not have voice related permissions
+    if channel.type is discord.ChannelType.text:
+        denied = discord.Permissions.voice()
+        base.value &= ~denied.value
+
+    return base
+
+
+def overwrite_from_dict(data):
+    allow = discord.Permissions(data.get("allow", 0))
+    deny = discord.Permissions(data.get("deny", 0))
+    return discord.PermissionOverwrite.from_pair(allow, deny)

punish/punish.py

@@ -1199,9 +1199,9 @@ class Punish(commands.Cog):
         if str(member.id) not in punished:
             if nitro_role != "no_nitro_role":
                 nitro_role = role_from_string(guild, nitro_role)
-                user_roles = {r for r in member.roles if r is not nitro_role}
-            else:
-                user_roles = {r for r in member.roles}
+                remove_role_set.discard(nitro_role)
+
+            user_roles = set(member.roles)
             # build lists of roles that *should* be removed and ones that *can* be
             removed_roles = user_roles & remove_role_set
             too_high_to_remove = {r for r in removed_roles if r >= guild.me.top_role}
@@ -1296,15 +1296,10 @@ class Punish(commands.Cog):
             await self.cancel_queue_event(member.guild.id, member.id)
 
             if apply_roles:
-                if nitro_role != "no_nitro_role":
-                    nitro_role = role_from_string(guild, nitro_role)
-                    user_roles = {r for r in member.roles if r is not nitro_role}
-                else:
-                    user_roles = {r for r in member.roles}
 
                 # readd removed roles from user, by replacing user's roles with all of their roles plus the ones that
                 # were removed (and can be re-added), minus the punish role
-                user_roles = {r for r in member.roles if r is not nitro_role}
+                user_roles = set(member.roles)
                 too_high_to_restore = {r for r in removed_roles if r >= guild.me.top_role}
                 removed_roles -= too_high_to_restore
                 user_roles |= removed_roles

punish/utils.py

@@ -179,12 +179,12 @@ def permissions_for_roles(channel, *roles):
     allows = 0
 
     # Apply channel specific role permission overwrites
-    for overwrite in channel.overwrites:
+    for target, overwrite in channel.overwrites.items():
         # Handle default role first, if present
         if overwrite.id == default.id:
             base.handle_overwrite(allow=overwrite.allow, deny=overwrite.deny)
 
-        if overwrite.type == "role" and overwrite.id in role_ids:
+        if isinstance(target, discord.Role) and target.id in role_ids:
             denies |= overwrite.deny
             allows |= overwrite.allow