import asyncio
import discord
from sqlitedict import SqliteDict
from services . environment_service import EnvService
from services . moderations_service import Moderation , ThresholdSet
MOD_DB = None
try :
print ( " Attempting to retrieve the General and Moderations DB " )
MOD_DB = SqliteDict ( " main_db.sqlite " , tablename = " moderations " , autocommit = True )
except Exception as e :
print ( " Failed to retrieve the General and Moderations DB " )
raise e
class ModerationsService ( discord . Cog , name = " ModerationsService " ) :
""" Cog containing moderation tools and features """
def __init__ (
self ,
bot ,
usage_service ,
model ,
) :
super ( ) . __init__ ( )
self . bot = bot
self . usage_service = usage_service
self . model = model
# Moderation service data
self . moderation_queues = { }
self . moderation_alerts_channel = EnvService . get_moderations_alert_channel ( )
self . moderation_enabled_guilds = [ ]
self . moderation_tasks = { }
self . moderations_launched = [ ]
# Defaults
self . default_warn_set = ThresholdSet ( 0.01 , 0.05 , 0.05 , 0.91 , 0.1 , 0.45 , 0.1 )
self . default_delete_set = ThresholdSet ( 0.26 , 0.26 , 0.1 , 0.95 , 0.03 , 0.85 , 0.4 )
@discord.Cog.listener ( )
async def on_ready ( self ) :
""" Check moderation service for each guild """
for guild in self . bot . guilds :
self . get_or_set_warn_set ( guild . id )
self . get_or_set_delete_set ( guild . id )
await self . check_and_launch_moderations ( guild . id )
print ( " The moderation service is ready. " )
def check_guild_moderated ( self , guild_id ) :
""" Given guild id, return bool of moderation status """
return guild_id in MOD_DB and MOD_DB [ guild_id ] [ " moderated " ]
def get_moderated_alert_channel ( self , guild_id ) :
""" Given guild id, return alert channel """
return MOD_DB [ guild_id ] [ " alert_channel " ]
def set_moderated_alert_channel ( self , guild_id , channel_id ) :
""" Given guild id and channel id, set channel to recieve alerts """
MOD_DB [ guild_id ] = { " moderated " : True , " alert_channel " : channel_id }
MOD_DB . commit ( )
def get_or_set_warn_set ( self , guild_id ) :
""" Get warn_set set for the guild, if not set them from default values """
guild_id = str ( guild_id )
key = guild_id + " _warn_set "
if key not in MOD_DB :
MOD_DB [ key ] = zip (
self . default_warn_set . keys , self . default_warn_set . thresholds
)
MOD_DB . commit ( )
return dict ( MOD_DB [ key ] )
def get_or_set_delete_set ( self , guild_id ) :
""" Get delete_set set for the guild, if not set them from default values """
guild_id = str ( guild_id )
key = guild_id + " _delete_set "
if key not in MOD_DB :
MOD_DB [ key ] = zip (
self . default_delete_set . keys , self . default_delete_set . thresholds
)
MOD_DB . commit ( )
return dict ( MOD_DB [ key ] )
def set_warn_set ( self , guild_id , threshold_set ) :
""" Set threshold for warning a message """
guild_id = str ( guild_id )
key = guild_id + " _warn_set "
MOD_DB [ key ] = zip ( threshold_set . keys , threshold_set . thresholds )
MOD_DB . commit ( )
def set_delete_set ( self , guild_id , threshold_set ) :
""" Set threshold for deleting a message """
guild_id = str ( guild_id )
key = guild_id + " _delete_set "
MOD_DB [ key ] = zip ( threshold_set . keys , threshold_set . thresholds )
MOD_DB . commit ( )
def set_guild_moderated ( self , guild_id , status = True ) :
""" Set the guild to moderated or not """
if guild_id not in MOD_DB :
MOD_DB [ guild_id ] = { " moderated " : status , " alert_channel " : 0 }
MOD_DB . commit ( )
return
MOD_DB [ guild_id ] = {
" moderated " : status ,
" alert_channel " : self . get_moderated_alert_channel ( guild_id ) ,
}
MOD_DB . commit ( )
async def check_and_launch_moderations ( self , guild_id , alert_channel_override = None ) :
""" Create the moderation service """
if self . check_guild_moderated ( guild_id ) :
Moderation . moderation_queues [ guild_id ] = asyncio . Queue ( )
moderations_channel = await self . bot . fetch_channel (
self . get_moderated_alert_channel ( guild_id )
if not alert_channel_override
else alert_channel_override
)
warn_set_nums = self . get_or_set_warn_set ( guild_id ) . values ( )
delete_set_nums = self . get_or_set_delete_set ( guild_id ) . values ( )
warn_set = ThresholdSet ( * warn_set_nums )
delete_set = ThresholdSet ( * delete_set_nums )
Moderation . moderation_tasks [ guild_id ] = asyncio . ensure_future (
Moderation . process_moderation_queue (
Moderation . moderation_queues [ guild_id ] ,
1 ,
1 ,
moderations_channel ,
warn_set ,
delete_set ,
)
)
print ( " Launched the moderations service for guild " + str ( guild_id ) )
Moderation . moderations_launched . append ( guild_id )
return moderations_channel
return None
async def moderations_command (
self , ctx : discord . ApplicationContext , status : str , alert_channel_id : str
) :
""" command handler for toggling moderation and setting an alert channel """
await ctx . defer ( )
try :
if alert_channel_id :
int ( alert_channel_id )
except ValueError :
# the alert_channel_id was passed in as a channel NAME instead of an ID, fetch the ID.
alert_channel = discord . utils . get ( ctx . guild . channels , name = alert_channel_id )
alert_channel_id = alert_channel . id
if status == " on " :
# Check if the current guild is already in the database and if so, if the moderations is on
if self . check_guild_moderated ( ctx . guild_id ) :
await ctx . respond ( " Moderations is already enabled for this guild " )
return
# Create the moderations service.
await self . start_moderations_service (
guild_id = ctx . guild_id , alert_channel_id = alert_channel_id
)
await ctx . respond ( " Moderations is now enabled for this guild " )
elif status == " off " :
# Cancel the moderations service.
await self . stop_moderations_service ( ctx . guild_id )
await ctx . respond (
" Moderations is now disabled for this guild " , ephemeral = True
)
async def stop_moderations_service ( self , guild_id ) :
""" Remove guild moderation status and stop the service """
self . set_guild_moderated ( guild_id , False )
Moderation . moderation_tasks [ guild_id ] . cancel ( )
Moderation . moderation_tasks [ guild_id ] = None
Moderation . moderation_queues [ guild_id ] = None
Moderation . moderations_launched . remove ( guild_id )
async def start_moderations_service ( self , guild_id , alert_channel_id = None ) :
""" Set guild moderation and start the service """
self . set_guild_moderated ( guild_id )
moderations_channel = await self . check_and_launch_moderations (
guild_id ,
Moderation . moderation_alerts_channel
if not alert_channel_id
else alert_channel_id ,
)
self . set_moderated_alert_channel ( guild_id , moderations_channel . id )
async def restart_moderations_service ( self , ctx ) :
""" restarts the moderation of the guild it ' s run in """
if not self . check_guild_moderated ( ctx . guild_id ) :
await ctx . respond (
" Moderations are not enabled, can ' t restart " ,
ephemeral = True ,
delete_after = 30 ,
)
return
await ctx . respond (
" The moderations service is being restarted... " ,
ephemeral = True ,
delete_after = 30 ,
)
await self . stop_moderations_service ( ctx . guild_id )
await ctx . send_followup (
" The moderations service was stopped.. " , ephemeral = True , delete_after = 30
)
await self . start_moderations_service (
ctx . guild_id , self . get_moderated_alert_channel ( ctx . guild_id )
)
await ctx . send_followup (
" The moderations service was restarted successfully. " ,
ephemeral = True ,
delete_after = 30 ,
)
async def build_moderation_settings_embed ( self , category , mod_set ) :
embed = discord . Embed (
title = " Moderation Settings " ,
description = " The moderation settings for this guild for the type: "
+ category ,
color = discord . Color . yellow ( ) if type == " warn " else discord . Color . red ( ) ,
)
# Add each key_value pair in the mod_set to the embed, make them fairly small
for key , value in mod_set . items ( ) :
embed . add_field ( name = key , value = value , inline = False )
return embed
async def config_command (
self ,
ctx : discord . ApplicationContext ,
config_type : str ,
hate ,
hate_threatening ,
self_harm ,
sexual ,
sexual_minors ,
violence ,
violence_graphic ,
) :
""" command handler for assigning threshold values for warn or delete """
all_args = [
hate ,
hate_threatening ,
self_harm ,
sexual ,
sexual_minors ,
violence ,
violence_graphic ,
]
await ctx . defer ( ephemeral = True )
# Case for printing the current config
if not any ( all_args ) :
await ctx . respond (
ephemeral = True ,
embed = await self . build_moderation_settings_embed (
config_type ,
self . get_or_set_warn_set ( ctx . guild_id )
if config_type == " warn "
else self . get_or_set_delete_set ( ctx . guild_id ) ,
) ,
)
return
if config_type == " warn " :
# Check if no args were
warn_set = self . get_or_set_warn_set ( ctx . guild_id )
new_warn_set = ThresholdSet (
hate if hate else warn_set [ " hate " ] ,
hate_threatening if hate_threatening else warn_set [ " hate/threatening " ] ,
self_harm if self_harm else warn_set [ " self-harm " ] ,
sexual if sexual else warn_set [ " sexual " ] ,
sexual_minors if sexual_minors else warn_set [ " sexual/minors " ] ,
violence if violence else warn_set [ " violence " ] ,
violence_graphic if violence_graphic else warn_set [ " violence/graphic " ] ,
)
self . set_warn_set ( ctx . guild_id , new_warn_set )
await self . restart_moderations_service ( ctx )
elif config_type == " delete " :
delete_set = self . get_or_set_delete_set ( ctx . guild_id )
new_delete_set = ThresholdSet (
hate if hate else delete_set [ " hate " ] ,
hate_threatening
if hate_threatening
else delete_set [ " hate/threatening " ] ,
self_harm if self_harm else delete_set [ " self-harm " ] ,
sexual if sexual else delete_set [ " sexual " ] ,
sexual_minors if sexual_minors else delete_set [ " sexual/minors " ] ,
violence if violence else delete_set [ " violence " ] ,
violence_graphic
if violence_graphic
else delete_set [ " violence/graphic " ] ,
)
self . set_delete_set ( ctx . guild_id , new_delete_set )
await self . restart_moderations_service ( ctx )
async def moderations_test_command (
self , ctx : discord . ApplicationContext , prompt : str
) :
""" command handler for checking moderation values of a given input """
await ctx . defer ( )
response = await self . model . send_moderations_request ( prompt )
await ctx . respond ( response [ " results " ] [ 0 ] [ " category_scores " ] )
await ctx . send_followup ( response [ " results " ] [ 0 ] [ " flagged " ] )