Made NSFW module open source again

.gitignore

@@ -5,7 +5,6 @@ exec
 *.out
 __pycache__
 *.opus
-nsfw.py
 unknown_errors.txt
 *.jpg
 token.py

wbb/modules/nsfw.py

@@ -0,0 +1,173 @@
+from os import remove
+
+from pyrogram import filters
+
+from wbb import app, arq
+from wbb.core.decorators.errors import capture_err
+from wbb.modules.admin import member_permissions
+from wbb.utils.dbfunctions import is_nsfw_on, nsfw_off, nsfw_on
+from wbb.utils.filter_groups import nsfw_detect_group
+
+__MODULE__ = "NSFW"
+__HELP__ = """
+/nsfw_scan - Manually Scan An Image/Sticker/Document.
+/anti_nsfw [ENABLE | DISABLE] - Turn This Module On/Off
+"""
+
+
+async def get_file_id_from_message(message):
+    file_id = None
+    if message.document:
+        if int(message.document.file_size) > 3145728:
+            return
+        mime_type = message.document.mime_type
+        if mime_type != "image/png" and mime_type != "image/jpeg":
+            return
+        file_id = message.document.file_id
+
+    if message.sticker:
+        if message.sticker.is_animated:
+            if not message.sticker.thumbs:
+                return
+            file_id = message.sticker.thumbs[0].file_id
+        else:
+            file_id = message.sticker.file_id
+
+    if message.photo:
+        file_id = message.photo.file_id
+
+    if message.animation:
+        if not message.animation.thumbs:
+            return
+        file_id = message.animation.thumbs[0].file_id
+
+    if message.video:
+        if not message.video.thumbs:
+            return
+        file_id = message.video.thumbs[0].file_id
+    return file_id
+
+
+@app.on_message(
+    (
+        filters.document
+        | filters.photo
+        | filters.sticker
+        | filters.animation
+        | filters.video
+    )
+    & ~filters.private,
+    group=nsfw_detect_group,
+)
+@capture_err
+async def detect_nsfw(_, message):
+    if not await is_nsfw_on(message.chat.id):
+        return
+    if not message.from_user:
+        return
+    file_id = await get_file_id_from_message(message)
+    if not file_id:
+        return
+    file = await app.download_media(file_id)
+    try:
+        results = await arq.nsfw_scan(file=file)
+    except Exception:
+        return
+    if not results.ok:
+        return
+    results = results.result
+    remove(file)
+    nsfw = results.is_nsfw
+    if not nsfw:
+        return
+    try:
+        await message.delete()
+    except Exception:
+        return
+    await message.reply_text(
+        f"""
+**NSFW Image Detected & Deleted Successfully!
+————————————————————————**
+**User:** {message.from_user.mention} [`{message.from_user.id}`]
+**Safe:** `{results.neutral} %`
+**Porn:** `{results.porn} %`
+**Adult:** `{results.sexy} %`
+**Hentai:** `{results.hentai} %`
+**Drawings:** `{results.drawings} %`
+**————————————————————————**
+__Use '/anti_nsfw disable' to disable this.__
+"""
+    )
+
+
+@app.on_message(filters.command("nsfw_scan"))
+@capture_err
+async def nsfw_scan_command(_, message):
+    if not message.reply_to_message:
+        await message.reply_text(
+            "Reply to an image/document/sticker/animation to scan it."
+        )
+        return
+    reply = message.reply_to_message
+    if (
+        not reply.document
+        and not reply.photo
+        and not reply.sticker
+        and not reply.animation
+        and not reply.video
+    ):
+        await message.reply_text(
+            "Reply to an image/document/sticker/animation to scan it."
+        )
+        return
+    m = await message.reply_text("Scanning")
+    file_id = await get_file_id_from_message(reply)
+    if not file_id:
+        return await m.edit("Something wrong happened.")
+    file = await app.download_media(file_id)
+    try:
+        results = await arq.nsfw_scan(file=file)
+    except Exception:
+        return
+    remove(file)
+    if not results.ok:
+        return await m.edit(results.result)
+    results = results.result
+    await m.edit(
+        f"""
+**Neutral:** `{results.neutral} %`
+**Porn:** `{results.porn} %`
+**Hentai:** `{results.hentai} %`
+**Sexy:** `{results.sexy} %`
+**Drawings:** `{results.drawings} %`
+**NSFW:** `{results.is_nsfw}`
+"""
+    )
+
+
+@app.on_message(filters.command("anti_nsfw") & ~filters.private)
+@capture_err
+async def nsfw_enable_disable(_, message):
+    if len(message.command) != 2:
+        await message.reply_text("Usage: /anti_nsfw [enable | disable]")
+        return
+    status = message.text.split(None, 1)[1].strip()
+    status = status.lower()
+    chat_id = message.chat.id
+    user_id = message.from_user.id
+    permissions = await member_permissions(chat_id, user_id)
+    if "can_change_info" not in permissions:
+        await message.reply_text("You don't have enough permissions.")
+        return
+    if status == "enable":
+        await nsfw_on(chat_id)
+        await message.reply_text(
+            "Enabled AntiNSFW System. I will Delete Messages Containing Inappropriate Content."
+        )
+    elif status == "disable":
+        await nsfw_off(chat_id)
+        await message.reply_text("Disabled AntiNSFW System.")
+    else:
+        await message.reply_text(
+            "Unknown Suffix, Use /anti_nsfw [enable|disable]"
+        )