505 lines
18 KiB
Python
505 lines
18 KiB
Python
import os
|
|
import glob
|
|
import sqlite3
|
|
import interactions
|
|
import yaml
|
|
import random
|
|
import string
|
|
import subprocess
|
|
from interactions import *
|
|
|
|
# Get current working directory
|
|
main_directory = os.getcwd()
|
|
|
|
# Create database and table for user and api keys
|
|
if not os.path.isfile(f"{main_directory}/database.db"):
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute(
|
|
'CREATE TABLE IF NOT EXISTS "DATABASE" ( "user" TEXT, "apikey" TEXT, "status" TEXT, PRIMARY KEY("user") )')
|
|
dbconnection.close()
|
|
|
|
# Making sure a serve.yml file exists and using that as the config file.
|
|
try:
|
|
serve = glob.glob(f'{main_directory}/serve.yml')[0]
|
|
except:
|
|
with open(f'{main_directory}/Instructions.txt', 'w') as Instructions:
|
|
print(f"Please create your serve.yml in {main_directory}")
|
|
Instructions.write("Place your serve.yml in this directory!")
|
|
exit()
|
|
|
|
# Check if discord bot token file exists, if not create file
|
|
if not os.path.isfile(f"{main_directory}/bot-token.txt"):
|
|
with open(f'{main_directory}/bot-token.txt', 'w') as discord_bot_token:
|
|
print("Please put your bot token in bot-token.txt")
|
|
discord_bot_token.write("Delete this and place your bot token on this line")
|
|
exit()
|
|
|
|
# Check if bot token exists and if it has been changed
|
|
with open(f'{main_directory}/bot-token.txt') as discord_bot_token:
|
|
bot_token = discord_bot_token.readline()
|
|
if bot_token == "Delete this and place your bot token on this line":
|
|
print("Please put your bot token in bot-token.txt")
|
|
exit()
|
|
|
|
bot = interactions.Client(token=bot_token)
|
|
|
|
|
|
# API Key generator function
|
|
async def generate_api_key():
|
|
key = "".join(random.choices(string.ascii_letters + string.digits, k=32))
|
|
return key
|
|
|
|
|
|
# Initial API Request function
|
|
async def initial_request(user: str):
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute("INSERT or REPLACE INTO database VALUES (?, ?, ?)", (user, "", "requested"))
|
|
dbconnection.commit()
|
|
dbconnection.close()
|
|
|
|
|
|
# Approve request function
|
|
async def approve_request(user: str):
|
|
# Open DB and approve the request, store the API key.
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
api_key = await generate_api_key()
|
|
dbcursor.execute("INSERT or REPLACE INTO database VALUES (?, ?, ?)", (user, api_key, "approved"))
|
|
dbconnection.commit()
|
|
dbconnection.close()
|
|
|
|
# Open YAML file
|
|
with open(serve) as serve_yaml:
|
|
yaml_dict = yaml.safe_load(serve_yaml)
|
|
|
|
# Add new entries
|
|
yaml_dict["users"][api_key] = {"username": user}
|
|
yaml_dict["users"][api_key]["devices"] = ["CDM"]
|
|
|
|
# Write the new YAML
|
|
with open(serve, "w") as new_serve_yaml:
|
|
yaml.safe_dump(yaml_dict, new_serve_yaml)
|
|
|
|
return api_key
|
|
|
|
|
|
# Deny user request function
|
|
async def deny_request(user: str):
|
|
# Open DB and deny the request.
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute("INSERT or REPLACE INTO database VALUES (?, ?, ?)", (user, "", "denied"))
|
|
dbconnection.commit()
|
|
dbconnection.close()
|
|
|
|
|
|
# Define check request status function
|
|
async def check_request_status(user: str):
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute("SELECT status FROM database WHERE user = :user", {"user": user})
|
|
user_status = dbcursor.fetchall()
|
|
formatted_user_status = ''
|
|
for character in str(user_status):
|
|
if character.isalnum():
|
|
formatted_user_status += character
|
|
dbconnection.close()
|
|
return formatted_user_status
|
|
|
|
|
|
# Define retrieve API key function
|
|
async def retrieve_api_key(user: str):
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute("SELECT apikey FROM database WHERE user = :user", {"user": user})
|
|
user_api_key = dbcursor.fetchall()
|
|
formatted_api_key = ''
|
|
for character in str(user_api_key):
|
|
if character.isalnum():
|
|
formatted_api_key += character
|
|
dbconnection.close()
|
|
return formatted_api_key
|
|
|
|
|
|
# Revoke user API key function
|
|
async def revoke_key(user: str):
|
|
# Find the users API key to remove from the serve.yml
|
|
user_api_key = await retrieve_api_key(user)
|
|
|
|
# Open DB and revoke the api key.
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute("INSERT or REPLACE INTO database VALUES (?, ?, ?)", (user, "", "revoked"))
|
|
dbconnection.commit()
|
|
dbconnection.close()
|
|
|
|
# Open the YAML file
|
|
with open(serve) as serve_yaml:
|
|
yaml_dict = yaml.safe_load(serve_yaml)
|
|
|
|
# Delete the API key entry
|
|
del yaml_dict["users"][user_api_key]
|
|
|
|
# Save the YAML file
|
|
with open(serve, "w") as new_serve_yaml:
|
|
yaml.safe_dump(yaml_dict, new_serve_yaml)
|
|
|
|
|
|
# Reset user status
|
|
async def reset_status(user: str):
|
|
# Open DB and reset the user status.
|
|
dbconnection = sqlite3.connect("database.db")
|
|
dbcursor = dbconnection.cursor()
|
|
dbcursor.execute("INSERT or REPLACE INTO database VALUES (?, ?, ?)", (user, "", ""))
|
|
dbconnection.commit()
|
|
dbconnection.close()
|
|
|
|
|
|
# Request API key slash command
|
|
@slash_command(name="request_api_key", description="Request an API key from https://api.cdm-project.com/")
|
|
async def request_api_key(ctx: SlashContext):
|
|
# Set admin channel
|
|
admin_channel = await bot.fetch_channel(1174897288745861120, force=True)
|
|
|
|
# Get user information and set DM
|
|
user_id = ctx.user.id
|
|
user_name = ctx.user.username
|
|
user_dm = bot.get_user(user_id)
|
|
|
|
# API request modal
|
|
api_request_modal = Modal(
|
|
ParagraphText(label="Why do you want an API key? (Be specific)", custom_id="request_reason"),
|
|
title="CDM API access",
|
|
)
|
|
|
|
# API request buttons
|
|
api_request_buttons: list[ActionRow] = [
|
|
ActionRow(
|
|
Button(
|
|
custom_id=f"{str(user_id)}Approve",
|
|
style=ButtonStyle.GREEN,
|
|
label="Approve",
|
|
),
|
|
Button(
|
|
custom_id=f"{str(user_id)}Deny",
|
|
style=ButtonStyle.RED,
|
|
label="Deny",
|
|
)
|
|
)
|
|
]
|
|
|
|
# API approved button
|
|
approved_button = Button(
|
|
custom_id=str(user_id),
|
|
style=ButtonStyle.BLUE,
|
|
label="Approved!",
|
|
disabled=True
|
|
)
|
|
|
|
# API denied button
|
|
denied_button = Button(
|
|
custom_id=str(user_id),
|
|
style=ButtonStyle.DANGER,
|
|
label="Denied!",
|
|
disabled=True
|
|
)
|
|
|
|
# Check if user has submitted an application before
|
|
|
|
# Send application and request if not already submitted
|
|
if await check_request_status(str(user_id)) != "requested" and await check_request_status(
|
|
str(user_id)) != "approved" and await check_request_status(str(user_id)) != "revoked":
|
|
await ctx.send_modal(modal=api_request_modal)
|
|
modal_ctx: ModalContext = await ctx.bot.wait_for_modal(api_request_modal)
|
|
await modal_ctx.send(f"{user_name}, Your application for an API key has been submitted.")
|
|
await initial_request(str(user_id))
|
|
admin_request_message = await admin_channel.send(
|
|
f"{user_name} has requested API access.\n```{modal_ctx.responses['request_reason']}```",
|
|
components=api_request_buttons)
|
|
status = await bot.wait_for_component(components=api_request_buttons)
|
|
|
|
# Check if the button was clicked for approve and send response
|
|
if status.ctx.custom_id == f"{str(user_id)}Approve":
|
|
await status.ctx.edit_origin(components=approved_button)
|
|
await admin_request_message.edit(content=f"Approved {user_name}!")
|
|
api_key = await approve_request(str(user_id))
|
|
api_restart_command = "sudo systemctl restart CDM-API.service"
|
|
subprocess.run(api_restart_command, shell=True)
|
|
await user_dm.send(f"Your request has been approved!\n\nYour API Key: `{api_key}`")
|
|
await ctx.member.add_role(1107788510129295510)
|
|
|
|
# Check if the button was clicked for deny and send response
|
|
elif status.ctx.custom_id == f"{str(user_id)}Deny":
|
|
await status.ctx.edit_origin(components=denied_button)
|
|
await admin_request_message.edit(content=f"Denied {user_name}!")
|
|
await deny_request(str(user_id))
|
|
await user_dm.send(f"You've been denied! Please either submit with a better reason or donate!")
|
|
|
|
# Send message if already submitted
|
|
else:
|
|
await ctx.send("Application already submitted, please use `/check_status` to view your application status")
|
|
|
|
|
|
# Check status slash command
|
|
@slash_command(name="check_status", description="Check application status / API key")
|
|
async def check_status(ctx: SlashContext):
|
|
# API key approved embed
|
|
api_key_approved_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Approved! ✔")
|
|
api_key_approved_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# API Key denied embed
|
|
api_key_denied_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Denied ✖")
|
|
api_key_denied_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# API Key denied embed
|
|
api_key_revoked_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Revoked! ✖")
|
|
api_key_denied_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# API key pending embed
|
|
api_key_pending_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Pending ❓")
|
|
api_key_pending_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# Get User ID from interaction
|
|
user_id = ctx.user.id
|
|
|
|
# Check status of API key request
|
|
status = await check_request_status(str(user_id))
|
|
|
|
# Handle response based on status
|
|
if str(status) == "requested":
|
|
await ctx.send(f"Your API status is `{str(status).strip('[](),')}`", embeds=api_key_pending_embed)
|
|
elif str(status) == "denied":
|
|
await ctx.send(f"Your API status is `{str(status).strip('[](),')}`", embeds=api_key_denied_embed)
|
|
elif str(status) == "revoked":
|
|
await ctx.send(f"Your API status is `{str(status).strip('[](),')}`", embeds=api_key_revoked_embed)
|
|
elif str(status) == "approved":
|
|
user_api_key = await retrieve_api_key(str(user_id))
|
|
api_key_approved_embed.add_field("API Key", user_api_key)
|
|
await ctx.send(f"Your API status is `{str(status).strip('[](),')}`", embeds=api_key_approved_embed,
|
|
ephemeral=True)
|
|
else:
|
|
await ctx.send("API permissions not requested yet, please use `/request_api`")
|
|
|
|
|
|
# Admin check status slash command
|
|
@slash_command(
|
|
name="check_user_status",
|
|
description="Check application status / API key for specific user",
|
|
default_member_permissions=Permissions.ADMINISTRATOR,
|
|
dm_permission=False
|
|
)
|
|
@slash_option(
|
|
name="discord_id",
|
|
description="Discord ID of the user you want to check.",
|
|
required=True,
|
|
opt_type=OptionType.STRING
|
|
)
|
|
async def check_user_status(ctx: SlashContext, discord_id: str):
|
|
# API key approved embed
|
|
api_key_approved_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Approved! ✔")
|
|
api_key_approved_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# API Key denied embed
|
|
api_key_denied_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Denied ✖")
|
|
api_key_denied_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# API Key denied embed
|
|
api_key_revoked_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Revoked! ✖")
|
|
api_key_denied_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# API key pending embed
|
|
api_key_pending_embed = interactions.Embed(
|
|
title=f"API Key request",
|
|
description=f"Pending ❓")
|
|
api_key_pending_embed.set_footer(text=f"https://api.cdm-project.com/")
|
|
|
|
# Check status of API key request
|
|
status = await check_request_status(discord_id)
|
|
|
|
# Handle response based on status
|
|
if str(status) == "requested":
|
|
await ctx.send(f"{discord_id} API status is `{str(status).strip('[](),')}`", embeds=api_key_pending_embed,
|
|
ephemeral=True)
|
|
elif str(status) == "denied":
|
|
await ctx.send(f"{discord_id} API status is `{str(status).strip('[](),')}`", embeds=api_key_denied_embed,
|
|
ephemeral=True)
|
|
elif str(status) == "revoked":
|
|
await ctx.send(f"{discord_id} API status is `{str(status).strip('[](),')}`", embeds=api_key_revoked_embed,
|
|
ephemeral=False)
|
|
elif str(status) == "approved":
|
|
user_api_key = await retrieve_api_key(discord_id)
|
|
api_key_approved_embed.add_field("API Key", user_api_key)
|
|
await ctx.send(f"{discord_id} API status is `{str(status).strip('[](),')}`", embeds=api_key_approved_embed,
|
|
ephemeral=True)
|
|
else:
|
|
await ctx.send(f"{discord_id} permissions not requested yet")
|
|
|
|
|
|
# Revoke API Key slash command
|
|
@slash_command(
|
|
name="revoke_api_key",
|
|
default_member_permissions=Permissions.ADMINISTRATOR,
|
|
dm_permission=False,
|
|
description="Revoke the users API key"
|
|
)
|
|
@slash_option(
|
|
name="discord_id",
|
|
description="Discord ID of the user you want to revoke.",
|
|
required=True,
|
|
opt_type=OptionType.STRING
|
|
)
|
|
async def revoke_api_key(ctx: SlashContext, discord_id: str):
|
|
# Retrieve the users API key
|
|
api_key = await retrieve_api_key(discord_id)
|
|
|
|
# API revoke buttons
|
|
api_revoke_buttons: list[ActionRow] = [
|
|
ActionRow(
|
|
Button(
|
|
custom_id=f"{discord_id}Revoke",
|
|
style=ButtonStyle.DANGER,
|
|
label="Revoke?",
|
|
disabled=False
|
|
),
|
|
Button(
|
|
custom_id=f"{discord_id}Cancel",
|
|
style=ButtonStyle.BLURPLE,
|
|
label="Cancel",
|
|
disabled=False
|
|
)
|
|
)
|
|
]
|
|
|
|
# API revoked button
|
|
revoked_button = Button(
|
|
custom_id=f"{discord_id}Revoked",
|
|
style=ButtonStyle.GREEN,
|
|
label="Revoked!",
|
|
disabled=True
|
|
)
|
|
|
|
# API revoked canceled button
|
|
canceled_button = Button(
|
|
custom_id=f"{discord_id}Canceled",
|
|
style=ButtonStyle.GRAY,
|
|
label="Canceled.",
|
|
disabled=True
|
|
)
|
|
|
|
# Revoke embed
|
|
revoke_api_key_embed = interactions.Embed(
|
|
title=f"Revoke API key?",
|
|
description=f"Revoke `{api_key}` ❓")
|
|
|
|
# Revoked embed
|
|
revoked_api_key_embed = interactions.Embed(
|
|
title=f"Revoke API key?",
|
|
description=f"Revoked `{api_key}` ✔")
|
|
|
|
# Send message to confirm revocation
|
|
revoke_message = await ctx.send(embeds=revoke_api_key_embed, components=api_revoke_buttons, ephemeral=False)
|
|
button_status = await bot.wait_for_component(components=api_revoke_buttons)
|
|
|
|
# Check if the button was clicked for revoke and send response
|
|
if button_status.ctx.custom_id == f"{discord_id}Revoke":
|
|
await revoke_key(discord_id)
|
|
api_restart_command = "sudo systemctl restart CDM-API.service"
|
|
subprocess.run(api_restart_command, shell=True)
|
|
await button_status.ctx.edit_origin(components=revoked_button)
|
|
await revoke_message.edit(embeds=revoked_api_key_embed)
|
|
|
|
# Check if the button was clicked for cancel and send response
|
|
elif button_status.ctx.custom_id == f"{discord_id}Cancel":
|
|
await button_status.ctx.edit_origin(components=canceled_button)
|
|
|
|
|
|
@slash_command(
|
|
name="reset_user_status",
|
|
default_member_permissions=Permissions.ADMINISTRATOR,
|
|
dm_permission=False,
|
|
description="Reset the users API status"
|
|
)
|
|
@slash_option(
|
|
name="discord_id",
|
|
description="Discord ID of the user you want to reset.",
|
|
required=True,
|
|
opt_type=OptionType.STRING
|
|
)
|
|
async def reset_user_status(ctx: SlashContext, discord_id: str):
|
|
# API reset buttons
|
|
api_reset_buttons: list[ActionRow] = [
|
|
ActionRow(
|
|
Button(
|
|
custom_id=f"{discord_id}Reset",
|
|
style=ButtonStyle.DANGER,
|
|
label="Reset?",
|
|
disabled=False
|
|
),
|
|
Button(
|
|
custom_id=f"{discord_id}Cancel",
|
|
style=ButtonStyle.BLURPLE,
|
|
label="Cancel",
|
|
disabled=False
|
|
)
|
|
)
|
|
]
|
|
|
|
# API reset confirm button
|
|
reset_button = Button(
|
|
custom_id=f"{discord_id}Restart",
|
|
style=ButtonStyle.GREEN,
|
|
label="Reset!",
|
|
disabled=True
|
|
)
|
|
|
|
# API reset canceled button
|
|
canceled_button = Button(
|
|
custom_id=f"{discord_id}Canceled",
|
|
style=ButtonStyle.GRAY,
|
|
label="Canceled.",
|
|
disabled=True
|
|
)
|
|
|
|
# Reset embed
|
|
reset_status_embed = interactions.Embed(
|
|
title=f"Reset User status?",
|
|
description=f"Reset `{discord_id}` ❓")
|
|
|
|
# Reset confirm embed
|
|
reset_status_confirm = interactions.Embed(
|
|
title=f"Reset User status?",
|
|
description=f"Reset `{discord_id}` ✔")
|
|
|
|
# Send message to confirm revocation
|
|
reset_message = await ctx.send(embeds=reset_status_embed, components=api_reset_buttons, ephemeral=False)
|
|
button_status = await bot.wait_for_component(components=api_reset_buttons)
|
|
|
|
# Check if the button was clicked for reset and send response
|
|
if button_status.ctx.custom_id == f"{discord_id}Reset":
|
|
await reset_status(discord_id)
|
|
await button_status.ctx.edit_origin(components=reset_button)
|
|
await reset_message.edit(embeds=reset_status_confirm)
|
|
|
|
# Check if the button was clicked for cancel and send response
|
|
elif button_status.ctx.custom_id == f"{discord_id}Cancel":
|
|
await button_status.ctx.edit_origin(components=canceled_button)
|
|
|
|
|
|
bot.start()
|