2020-02-01 12:49:44 +13:00
from redbot . core . utils . chat_formatting import *
from redbot . core import Config , checks , commands , bank
import discord
2020-04-22 04:19:44 +12:00
import asyncio
2020-02-01 13:37:32 +13:00
2020-05-05 10:03:39 +12:00
2020-02-05 18:13:23 +13:00
class PoorError ( commands . CheckFailure ) :
2020-02-01 12:49:44 +13:00
pass
2020-02-01 13:37:32 +13:00
2020-02-01 12:49:44 +13:00
# 0 -> member, 1 -> guild
NEW_RECEIPT_MESSAGE = " Hello {0.mention} ! This is will be your receipt for commands you pay for in {1.name} . \n The message below will auto update as you pay for commands. I will pin the message for you so you can refer back to this to track your spending. "
# 0 -> recepit number, 1 -> command name, 2->command cost, 3 -> currency name
RECEIPT_MESSAGE = " {0} . {1} : {2} {3} \n "
MAX_MSG_LEN = 2000
2020-02-01 13:37:32 +13:00
2020-02-01 12:49:44 +13:00
class CostManager ( commands . Cog ) :
"""
Allows customizing costs for any commands loaded into read ,
and also set users who are exempt from costs on a per member or per role level .
Hierarchy : user > role > guild_role
"""
2020-02-01 13:37:32 +13:00
2020-02-01 12:49:44 +13:00
def __init__ ( self , bot ) :
self . bot = bot
self . config = Config . get_conf ( self , identifier = 13291493293 , force_registration = True )
# commands: {
# command_name: {
# cost: int
# user_ids: {user_id:cost}
# role_ids: {role_id:cost} # cost can be 0 to make free
# },
# ...
# }
2020-02-01 13:37:32 +13:00
default_guild = { " FREE_ROLES " : [ ] , " COMMANDS " : { } }
2020-02-01 12:49:44 +13:00
self . config . register_guild ( * * default_guild )
self . config . register_member ( receipt = 0 )
self . bot . before_invoke ( self . cost_checker )
def cog_unload ( self ) :
self . bot . remove_before_invoke_hook ( self . cost_checker )
# permission hook checker for cost of command
async def cost_checker ( self , ctx ) :
cost = await self . get_cost ( ctx )
if cost == 0 :
return None
try :
await bank . withdraw_credits ( ctx . author , cost )
await self . update_receipt ( ctx , cost )
except ValueError :
raise PoorError ( f " member: { ctx . author . id } , guild: { ctx . guild . id } " )
async def get_cost ( self , ctx , member = None , command = None ) :
"""
Get cost of a command , respecting hierarchy
"""
2020-02-10 01:51:38 +13:00
if isinstance ( ctx . channel , discord . DMChannel ) :
return 0
2020-02-01 12:49:44 +13:00
guild = ctx . guild
if not member :
member = ctx . author
member_roles = { r . id for r in member . roles if r . name != " @everyone " }
if not command :
command = ctx . command . name
guild_data = await self . config . guild ( guild ) . all ( )
command_data = guild_data [ " COMMANDS " ]
# check if settings for command
if command not in command_data . keys ( ) :
return 0
command_data = command_data [ command ]
2020-02-09 22:06:08 +13:00
charged_roles = set ( command_data . get ( " role_ids " , { } ) . keys ( ) )
2020-02-01 12:49:44 +13:00
found_roles = charged_roles & member_roles
cost = 0
# check user cost
2020-02-09 22:06:08 +13:00
if str ( member . id ) in command_data . get ( " user_ids " , { } ) . keys ( ) :
2020-02-01 12:49:44 +13:00
cost = command_data [ " user_ids " ] [ str ( member . id ) ]
# check role cost, choose lowest cost if mutliple roles found.
elif found_roles :
cost = min ( [ command_data [ " role_ids " ] [ r ] for r in found_roles ] )
2020-02-01 13:37:32 +13:00
else : # get normal cost for command and check guild free roles
2020-02-01 12:49:44 +13:00
cost = command_data [ " cost " ]
found_roles = set ( guild_data [ " FREE_ROLES " ] ) & member_roles
if found_roles :
cost = 0
return cost
async def update_receipt ( self , ctx , cost ) :
"""
Update user ' s receipt, or make a new one if its new receipt or old one is somehow missing.
"""
guild = ctx . guild
member = ctx . author
command = ctx . command . name
receipt = await self . config . member ( member ) . receipt ( )
channel = member . dm_channel
if not channel :
await member . create_dm ( )
channel = member . dm_channel
msg_receipt = None
if receipt > 0 :
try :
msg_receipt = await channel . fetch_message ( receipt )
except ( Forbidden , HTTPException ) :
return
except NotFound :
pass
if not msg_receipt :
try :
await channel . send ( NEW_RECEIPT_MESSAGE . format ( member , guild ) )
except :
2020-02-01 13:37:32 +13:00
if receipt != - 1 : # send mention if first time getting this message
await ctx . send (
f " Hey { member . mention } , please turn on DMs from server members in your settings so I can send you receipts for purchase of commands. "
)
2020-02-01 12:49:44 +13:00
await self . config . member ( member ) . receipt . set ( - 1 )
return
currency_name = await bank . get_currency_name ( guild )
msg_receipt = await channel . send ( RECEIPT_MESSAGE . format ( 1 , command , cost , currency_name ) )
await msg_receipt . pin ( )
await self . config . member ( member ) . receipt . set ( msg_receipt . id )
return
# msg already found, so update it
currency_name = await bank . get_currency_name ( guild )
content = msg_receipt . content . split ( " \n " )
# TODO: this is really messy but ill fix it...
last_num = int ( content [ - 1 ] . split ( " . " ) [ 0 ] )
content . append ( RECEIPT_MESSAGE . format ( last_num + 1 , command , cost , currency_name ) )
2020-04-25 15:50:17 +12:00
while len ( " \n " . join ( content ) ) > MAX_MSG_LEN :
2020-02-01 12:49:44 +13:00
del content [ 0 ]
content = " \n " . join ( content )
await msg_receipt . edit ( content = content )
async def clean_data ( self , guild ) :
async with self . config . guild ( guild ) . FREE_ROLES ( ) as free :
roles = [ guild . get_role ( id ) for id in free ]
free = [ r . id for r in roles if r is not None ]
async with self . config . guild ( guild ) . COMMANDS ( ) as c :
for command_name in list ( c . keys ( ) ) :
data = c [ command_name ]
for user_id in data . get ( " user_ids " , { } ) . keys ( ) :
if not guild . get_member ( int ( user_id ) ) :
del c [ command_name ] [ " user_ids " ] [ user_id ]
for role_id in data . get ( " role_ids " , { } ) . keys ( ) :
if not guild . get_role ( int ( role_id ) ) :
del c [ command_name ] [ " role_ids " ] [ role_id ]
# check validity of arguments
def arg_check ( self , cost : int , command : str ) :
if not self . bot . get_command ( command ) or cost < 0 :
return False
return True
# format list of items
@staticmethod
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 " "
@commands.group ( name = " costset " , invoke_without_command = True )
@commands.guild_only ( )
@checks.admin ( )
async def costset ( self , ctx , cost : int , * , command_name : str = None ) :
"""
Sets and manage cost / bypasses of commands .
"""
if ctx . invoked_subcommand :
return
if not self . arg_check ( cost , command_name ) :
await ctx . send ( " Invalid cost or command! " )
return
async with self . config . guild ( ctx . guild ) . COMMANDS ( ) as c :
if not command_name in c . keys ( ) :
c [ command_name ] = { }
c [ command_name ] [ " cost " ] = cost
await ctx . tick ( )
@costset.command ( name = " role " )
async def cost_set_role ( self , ctx , command_name : str , cost : int , * , role : discord . Role ) :
"""
Set cost of command for specific role .
Set to 0 to make command free for role .
"""
if not self . arg_check ( cost , command_name ) :
await ctx . send ( " Invalid cost or command! " )
return
async with self . config . guild ( ctx . guild ) . COMMANDS ( ) as c :
if not command_name in c . keys ( ) :
c [ command_name ] = { }
if not " role_ids " in c [ command_name ] . keys ( ) :
c [ command_name ] [ " role_ids " ] = { }
c [ command_name ] [ " role_ids " ] [ str ( role . id ) ] = cost
await ctx . tick ( )
@costset.command ( name = " user " )
async def cost_set_user ( self , ctx , command_name : str , cost : int , * , member : discord . Member ) :
"""
Set cost of command for specific user .
Set to 0 to make command free for user .
"""
if not self . arg_check ( cost , command_name ) :
await ctx . send ( " Invalid cost or command! " )
return
async with self . config . guild ( ctx . guild ) . COMMANDS ( ) as c :
if not command_name in c . keys ( ) :
c [ command_name ] = { }
if not " user_ids " in c [ command_name ] . keys ( ) :
c [ command_name ] [ " user_ids " ] = { }
c [ command_name ] [ " user_ids " ] [ str ( member . id ) ] = cost
await ctx . tick ( )
@costset.command ( name = " clear " )
async def cost_set_clear ( self , ctx ) :
"""
Clear missing roles / members in cost config .
"""
await self . clean_data ( ctx . guild )
await ctx . tick ( )
@costset.command ( name = " owner-clear " )
@checks.is_owner ( )
async def cost_set_owner_clear ( self , ctx ) :
"""
Clear missing roles / members in every cost config .
"""
for guild in self . bot . guilds :
await self . clean_data ( guild )
await ctx . tick ( )
@costset.command ( name = " free-roles " )
async def cost_set_free_roles ( self , ctx , * , role_list : str = None ) :
"""
Set roles who can use all commands for free in server .
* * Note * * : This only applies to commands whose cost was set with this cog .
Role list should be a list of one or more * * role names or ids * * seperated by commas .
Roles in role list will be removed if already in the free role list , or added if they are not .
Role names are case sensitive !
Don ' t pass a role list to see the current roles
"""
if not role_list :
curr = await self . config . guild ( ctx . guild ) . FREE_ROLES ( )
if not curr :
await ctx . send ( " No roles defined. " )
else :
curr = [ ctx . guild . get_role ( role_id ) for role_id in curr ]
not_found = len ( [ r for r in curr if r is None ] )
curr = [ r . name for r in curr if curr is not None ]
if not_found :
2020-02-01 13:37:32 +13:00
await ctx . send (
f " { not_found } roles weren ' t found, please run { ctx . prefix } costset clear to remove these roles. \n Free Roles: { self . format_list ( * curr ) } "
)
2020-02-01 12:49:44 +13:00
else :
await ctx . send ( f " Free Roles: { self . format_list ( * curr ) } " )
return
role_list = role_list . strip ( ) . split ( " , " )
role_list = [ r . strip ( ) for r in role_list ]
not_found = set ( )
found = set ( )
added = set ( )
removed = set ( )
for role_name in role_list :
role = discord . utils . find ( lambda r : r . name == role_name , ctx . guild . roles )
# if couldnt find by role name, try to find by role id
if role is None :
role = discord . utils . find ( lambda r : r . id == role_name , ctx . guild . roles )
if role is None :
not_found . add ( role_name )
continue
found . add ( role )
if not_found :
2020-02-01 13:37:32 +13:00
await ctx . send (
warning ( " These roles weren ' t found, please try again: {} " . format ( self . format_list ( * not_found ) ) )
)
2020-02-01 12:49:44 +13:00
return
async with self . config . guild ( ctx . guild ) . FREE_ROLES ( ) as free :
for role in found :
if role . id in free :
free . remove ( role . id )
removed . add ( role . name )
else :
free . append ( role . id )
added . add ( role . name )
msg = " "
if added :
msg + = " Added: {} \n " . format ( self . format_list ( * added ) )
if removed :
msg + = " Removed: {} " . format ( self . format_list ( * removed ) )
await ctx . send ( msg )
@costset.command ( name = " list " )
2020-02-09 13:06:23 +13:00
async def cost_set_list ( self , ctx ) :
2020-02-01 12:49:44 +13:00
"""
List current cost settings for the guild
"""
guild = ctx . guild
guild_data = await self . config . guild ( guild ) . all ( )
free_roles = [ guild . get_role ( r ) for r in guild_data [ " FREE_ROLES " ] ]
free_roles = self . format_list ( * [ r . name for r in free_roles if r is not None ] )
commands = guild_data [ " COMMANDS " ]
msg = f " Guild Free Roles: { free_roles } \n \n Commands: \n "
for command_name , data in commands . items ( ) :
msg + = f " \t { command_name } : { data [ ' cost ' ] } \n "
msg + = " \t \t Role Costs: \n "
2020-02-09 22:06:08 +13:00
for role_id , cost in data . get ( " role_ids " , { } ) . items ( ) :
2020-02-01 12:49:44 +13:00
role = guild . get_role ( int ( role_id ) )
if not role :
continue
msg + = f " \t \t \t { role . name } : { cost } \n "
msg + = " \t \t User Costs: \n "
2020-02-09 22:06:08 +13:00
for user_id , cost in data . get ( " user_ids " , { } ) . items ( ) :
2020-02-01 12:49:44 +13:00
user = guild . get_member ( int ( user_id ) )
if not user :
continue
msg + = f " \t \t \t { user . name } : { cost } \n "
msg = pagify ( msg )
for m in msg :
await ctx . send ( box ( m , lang = " python " ) )
@commands.command ( name = " cost " )
2020-02-02 16:45:15 +13:00
@commands.guild_only ( )
2020-02-01 12:49:44 +13:00
async def get_cost_command ( self , ctx , command : str ) :
"""
Get cost of a command .
"""
if not self . arg_check ( 0 , command ) :
await ctx . send ( warning ( f " Command ` { command } ` not found! " ) )
return
cost = await self . get_cost ( ctx , command = command )
if cost == 0 :
await ctx . send ( f " { command } is free for you! " )
return
else :
currency_name = await bank . get_currency_name ( ctx . guild )
await ctx . send ( f " { command } costs ` { cost } ` { currency_name } for you. " )
@commands.Cog.listener ( )
async def on_member_remove ( self , member ) :
await self . clean_data ( member . guild )
@commands.Cog.listener ( )
async def on_guild_role_delete ( self , role ) :
await self . clean_data ( role . guild )
# Listens for poorerror
@commands.Cog.listener ( )
async def on_command_error ( self , ctx , exception ) :
if isinstance ( exception , PoorError ) :
cost = await self . get_cost ( ctx )
currency_name = await bank . get_currency_name ( ctx . guild )
balance = await bank . get_balance ( ctx . author )
2020-04-22 04:19:44 +12:00
message = await ctx . send (
2020-02-01 13:37:32 +13:00
f " Sorry { ctx . author . name } , you do not have enough { currency_name } to use that command. (Cost: { cost } , Balance: { balance } ) "
)
2020-04-22 04:19:44 +12:00
await asyncio . sleep ( 10 )
await message . delete ( )