Fox-V3/nudity/nudity.py

import pathlib

import discord
from nudenet import NudeClassifier
from redbot.core import Config, commands
from redbot.core.bot import Red
from redbot.core.data_manager import cog_data_path


class Nudity(commands.Cog):
    """
    V3 Cog Template
    """

    def __init__(self, bot: Red):
        super().__init__()
        self.bot = bot
        self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True)

        default_guild = {"enabled": False, "channel_id": None}

        self.config.register_guild(**default_guild)

        # self.detector = NudeDetector()
        self.classifier = NudeClassifier()

        self.data_path: pathlib.Path = cog_data_path(self)

        self.current_processes = 0

    async def red_delete_data_for_user(self, **kwargs):
        """Nothing to delete"""
        return

    @commands.command(aliases=["togglenudity"], name="nudity")
    async def nudity(self, ctx: commands.Context):
        """Toggle nude-checking on or off"""
        is_on = await self.config.guild(ctx.guild).enabled()
        await self.config.guild(ctx.guild).enabled.set(not is_on)
        await ctx.send("Nude checking is now set to {}".format(not is_on))

    @commands.command()
    async def nsfwchannel(self, ctx: commands.Context, channel: discord.TextChannel = None):
        if channel is None:
            await self.config.guild(ctx.guild).channel_id.set(None)
            await ctx.send("NSFW Channel cleared")
        else:
            if not channel.is_nsfw():
                await ctx.send("This channel isn't NSFW!")
                return
            else:
                await self.config.guild(ctx.guild).channel_id.set(channel.id)
                await ctx.send("NSFW channel has been set to {}".format(channel.mention))

    async def get_nsfw_channel(self, guild: discord.Guild):
        channel_id = await self.config.guild(guild).channel_id()

        if channel_id is None:
            return None
        else:
            return guild.get_channel(channel_id=channel_id)

    async def nsfw(self, message: discord.Message, images: dict):
        content = message.content
        guild: discord.Guild = message.guild
        if not content:
            content = "*`None`*"
        try:
            await message.delete()
        except discord.Forbidden:
            await message.channel.send("NSFW Image detected!")
            return

        embed = discord.Embed(title="NSFW Image Detected")
        embed.add_field(name="Original Message", value=content)
        embed.set_author(name=message.author.name, icon_url=message.author.avatar_url)
        await message.channel.send(embed=embed)

        nsfw_channel = await self.get_nsfw_channel(guild)

        if nsfw_channel is None:
            return
        else:
            for image, r in images.items():
                if r["unsafe"] > 0.7:
                    await nsfw_channel.send(
                        "NSFW Image from {}".format(message.channel.mention),
                        file=discord.File(image,),
                    )

    @commands.Cog.listener()
    async def on_message(self, message: discord.Message):
        is_private = isinstance(message.channel, discord.abc.PrivateChannel)

        if not message.attachments or is_private or message.author.bot:
            # print("did not qualify")
            return

        if await self.bot.cog_disabled_in_guild(self, message.guild):
            return

        try:
            is_on = await self.config.guild(message.guild).enabled()
        except AttributeError:
            return

        if not is_on:
            print("Not on")
            return

        channel: discord.TextChannel = message.channel

        if channel.is_nsfw():
            print("nsfw channel is okay")
            return

        check_list = []
        for attachment in message.attachments:
            # async with aiohttp.ClientSession() as session:
            #     img = await fetch_img(session, attachment.url)

            ext = attachment.filename

            temp_name = self.data_path / f"nudecheck{self.current_processes}_{ext}"

            self.current_processes += 1

            print("Pre attachment save")
            await attachment.save(temp_name)
            check_list.append(temp_name)

        print("Pre nude check")
        # nude_results = self.detector.detect(temp_name)
        nude_results = self.classifier.classify([str(n) for n in check_list])
        # print(nude_results)

        if True in [r["unsafe"] > 0.7 for r in nude_results.values()]:
            # print("Is nude")
            await message.add_reaction("❌")
            await self.nsfw(message, nude_results)
        else:
            # print("Is not nude")
            await message.add_reaction("✅")


# async def fetch_img(session, url):
#     with aiohttp.Timeout(10):
#         async with session.get(url) as response:
#             assert response.status == 200
#             return await response.read()