Browse Source

moved danbooru clone interaction to a function in utils instead of in the NSFW cog so it can be used in many commands.

tags/v2.1.0
Roxie Gibson 5 years ago
parent
commit
17b0530df8
2 changed files with 82 additions and 42 deletions
  1. +23
    -42
      roxbot/cogs/nsfw.py
  2. +59
    -0
      roxbot/utils.py

+ 23
- 42
roxbot/cogs/nsfw.py View File

# SOFTWARE. # SOFTWARE.




import random

import discord import discord
from discord.ext import commands from discord.ext import commands


return blacklist return blacklist




class NFSW():
class NFSW:
"""The NSFW cog is a collection of commands that post images from popular NSFW sites. """ """The NSFW cog is a collection of commands that post images from popular NSFW sites. """
def __init__(self, bot_client): def __init__(self, bot_client):
self.bot = bot_client self.bot = bot_client
} }
} }


@roxbot.checks.is_nsfw()
@commands.command(hidden=True)
async def gelbooru_clone(self, ctx, base_url, post_url, tags):
limit = 150
async def gelbooru_clone(self, ctx, base_url, endpoint_url, tags):
if isinstance(ctx.channel, discord.TextChannel): if isinstance(ctx.channel, discord.TextChannel):
tags = tags + tag_blacklist(ctx.guild)
page = random.randrange(20)
url = base_url + tags + '&limit=' + str(limit) + '%pid=' + str(page)
if isinstance(ctx.channel, discord.DMChannel):
cache_id = ctx.author.id
banned_tags = tag_blacklist(ctx.guild)
else: else:
cache_id = ctx.guild.id
# IF ID is not in cache, create cache for ID
if not self.cache.get(cache_id, False):
self.cache[cache_id] = []

posts = await roxbot.http.api_request(url)

if posts is None:
banned_tags = ""

post = await roxbot.utils.danbooru_clone_api_req(
ctx.channel,
base_url,
endpoint_url,
tags=tags,
banned_tags=banned_tags,
cache=self.cache
)

if not post:
return await ctx.send("Nothing was found. *psst, check the tags you gave me.*") return await ctx.send("Nothing was found. *psst, check the tags you gave me.*")

post = None
counter = 0
while counter < 20:
post = random.choice(posts)
md5 = post.get("md5") or post.get("hash")
if md5 not in self.cache[cache_id]:
self.cache[cache_id].append(md5)
if len(self.cache[cache_id]) > 10:
self.cache[cache_id].pop(0)
break
counter += 1

url = post.get("file_url")
if not url:
url = post_url + "{0[directory]}/{0[image]}".format(post)
output = await ctx.send(url)
else:
output = await ctx.send(post)
await roxbot.utils.delete_option(self.bot, ctx, output, self.bot.get_emoji(444410658101002261) or "❌") await roxbot.utils.delete_option(self.bot, ctx, output, self.bot.get_emoji(444410658101002261) or "❌")


@roxbot.checks.is_nsfw() @roxbot.checks.is_nsfw()
;e621 test ;e621 test
""" """
base_url = "https://e621.net/post/index.json?tags=" base_url = "https://e621.net/post/index.json?tags="
return await ctx.invoke(self.gelbooru_clone, base_url=base_url, post_url="", tags=tags)
return await self.gelbooru_clone(ctx, base_url, "", tags)


@roxbot.checks.is_nsfw() @roxbot.checks.is_nsfw()
@commands.command() @commands.command()
;rule34 test ;rule34 test
""" """
base_url = "https://rule34.xxx/index.php?page=dapi&s=post&q=index&json=1&tags=" base_url = "https://rule34.xxx/index.php?page=dapi&s=post&q=index&json=1&tags="
post_url = "https://img.rule34.xxx/images/"
return await ctx.invoke(self.gelbooru_clone, base_url=base_url, post_url=post_url, tags=tags)
endpoint_url = "https://img.rule34.xxx/images/"
return await self.gelbooru_clone(ctx, base_url, endpoint_url, tags)


@roxbot.checks.is_nsfw() @roxbot.checks.is_nsfw()
@commands.command() @commands.command()
;gelbooru test ;gelbooru test
""" """
base_url = "https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&tags=" base_url = "https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&tags="
post_url = "https://simg3.gelbooru.com/images/"
return await ctx.invoke(self.gelbooru_clone, base_url=base_url, post_url=post_url, tags=tags)
endpoint_url = "https://simg3.gelbooru.com/images/"
return await self.gelbooru_clone(ctx, base_url, endpoint_url, tags)


@commands.guild_only() @commands.guild_only()
@commands.has_permissions(manage_channels=True) @commands.has_permissions(manage_channels=True)
return await ctx.send("No valid option given.") return await ctx.send("No valid option given.")
return settings.update(nsfw, "nsfw") return settings.update(nsfw, "nsfw")



def setup(bot_client): def setup(bot_client):
bot_client.add_cog(NFSW(bot_client)) bot_client.add_cog(NFSW(bot_client))

+ 59
- 0
roxbot/utils.py View File

# SOFTWARE. # SOFTWARE.




import random
import asyncio import asyncio
import discord import discord
import argparse import argparse


from roxbot import http
from roxbot import guild_settings from roxbot import guild_settings
from roxbot.enums import EmbedColours from roxbot.enums import EmbedColours


for key, value in kwargs.items(): for key, value in kwargs.items():
embed.add_field(name=key, value=value) embed.add_field(name=key, value=value)
return await channel.send(embed=embed) return await channel.send(embed=embed)


async def danbooru_clone_api_req(channel, base_url, endpoint_url, cache=None, tags="", banned_tags=""):
"""Utility function that deals with danbooru clone api interaction.
It also deals with cache management for these interactions.

Params
=======
channel: discord.Channel
Channel command has been invoked in
base_url: str
Base url of the site
endpoint_url: str
Endpoint of images in the API. This is used if the API does not give this in its response.
cache: dict (optional)
Post cache. Were channel ID's are keys with values that are lists of identifiable info.
Cache is handled in this function and will be updated so that other functions can access it.
tags: str (optional)
tags to use in the search. Separated by spaces.
banned_tags: str (optional)
banned tags to append to the search. Separated by spaces with a - in front to remove them from search results.
"""
limit = "150"
tags = tags + banned_tags
page_number = str(random.randrange(20))
url = base_url + tags + '&limit=' + limit + '%pid=' + page_number

if isinstance(channel, discord.DMChannel):
cache_id = channel.id
else:
cache_id = channel.guild.id

# IF ID is not in cache, create cache for ID
if not cache.get(cache_id, False):
cache[cache_id] = []

posts = await http.api_request(url)

if posts is None:
return None

post = None
counter = 0
while counter < 20:
post = random.choice(posts)
md5 = post.get("md5") or post.get("hash")
if md5 not in cache[cache_id]:
cache[cache_id].append(md5)
if len(cache[cache_id]) > 10:
cache[cache_id].pop(0)
break
counter += 1

url = post.get("file_url")
if not url:
url = endpoint_url + "{0[directory]}/{0[image]}".format(post)
return url

Loading…
Cancel
Save