disarchive/config.py

import jsonschema
import os
import re
import yaml
import sys
import discord
import logging
from colorlog import ColoredFormatter

log_level = logging.DEBUG
log_format = (
    "  %(log_color)s%(levelname)-8s%(reset)s |"
    " %(log_color)s%(message)s%(reset)s"
)

logging.root.setLevel(log_level)
formatter = ColoredFormatter(log_format)

stream = logging.StreamHandler()
stream.setLevel(log_level)
stream.setFormatter(formatter)

LOG = logging.getLogger("pythonConfig")
LOG.setLevel(log_level)
LOG.addHandler(stream)

TOKEN = None
BOT_COLOR = None

DB_NAME = None
DB_ENGINE = None
DB_HOST = None
DB_PORT = None
DB_USER = None
DB_PASSWORD = None


schema = {
    "type": "object",
    "properties": {
        "bot_info": {
            "type": "object",
            "properties": {
                "token": {"type": "string"},
                "bot_color": {"type": "string", "default": "#fc5f4e"},
            },
            "required": ["token"],
        },
        "database": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "default": "disarchive"},
                "engine": {"type": "string"},
                "host": {"type": "string"},
                "port": {"type": "integer"},
                "user": {"type": "string"},
                "password": {"type": "string"},
            },
            "required": [
                "name",
                "engine",
                "host",
                "port",
                "user",
                "password",
            ],
        },
    },
    "required": ["bot_info", "database"],
}


# Attempt to load the config file, otherwise create a new template
def load_config():
    if os.path.exists("/.dockerenv"):
        file_path = "/config/config.yaml"
    else:
        file_path = "config.yaml"

    try:
        with open(file_path, "r") as f:
            file_contents = f.read()
            validate_config(file_contents)

    except FileNotFoundError:
        # Create a new config.yaml file with the template
        with open(file_path, "w") as f:
            f.write(
                """
bot_info:
    token:
    bot_color:

database:
    name:
    engine:
    host:
    port:
    user:
    password: """
            )

        sys.exit(
            LOG.critical(
                "Configuration file `config.yaml` has been generated. Please"
                " fill out all of the necessary information. Refer to the docs"
                " for information on what a specific configuration option is."
            )
        )


# Thouroughly validate all of the options in the config.yaml file
def validate_config(file_contents):
    global TOKEN, BOT_COLOR, DB_NAME, DB_ENGINE, DB_HOST, DB_PORT, DB_USER, DB_PASSWORD
    config = yaml.safe_load(file_contents)

    try:
        jsonschema.validate(config, schema)
    except jsonschema.ValidationError as e:
        sys.exit(LOG.critical(f"Error in config.yaml file: {e.message}"))

    # Make sure "bot_color" is a valid hex color
    hex_pattern_one = "^#([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$"
    hex_pattern_two = "^([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$"

    if "bot_color" in config["bot_info"]:
        if not bool(
            re.match(hex_pattern_one, config["bot_info"]["bot_color"])
        ) and not bool(
            re.match(hex_pattern_two, config["bot_info"]["bot_color"])
        ):
            LOG.critical(
                "Error in config.yaml file: bot_color is not a valid hex color"
            )
        else:
            BOT_COLOR = discord.Color(
                int((config["bot_info"]["bot_color"]).replace("#", ""), 16)
            )

    if config["database"]["engine"] not in [
        "sqlite",
        "mysql",
        "postgresql",
    ]:
        LOG.error(
            "database_engine must be either 'sqlite', 'mysql', or 'postgresql'"
        )
        return False
    else:
        DB_ENGINE = config["database"]["engine"]

    DB_NAME = config["database"]["name"]
    DB_HOST = config["database"]["host"]
    DB_PORT = config["database"]["port"]
    DB_USER = config["database"]["user"]
    DB_PASSWORD = config["database"]["password"]

    TOKEN = config["bot_info"]["token"]


"""
Template for embeds
"""


def create_embed(
    title: str = None,
    description: str = None,
    color=None,
    footer=None,
    thumbnail=None,
):
    embed = discord.Embed(
        title=title,
        description=description,
        color=color if color else BOT_COLOR,
    )

    if footer:
        embed.set_footer(text=footer)
    # else:
    #     embed.set_footer(
    #         text=datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") + " UTC"
    #     )

    if thumbnail:
        embed.set_thumbnail(url=thumbnail)

    return embed