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
NAMING_SCHEME = None
BOT_COLOR = None

SQLITE_NAME = "disarchive"
DB_NAME = None
DB_ENGINE = None
DB_HOST = None
DB_PORT = None
DB_USER = None
DB_PASSWORD = None

schema = {
    "type": "object",
    "properties": {
        "general": {
            "type": "object",
            "properties": {
                "token": {"type": "string"},
                "bot_color": {"type": "string", "default": "#fc5f4e"},
                "naming_scheme": {
                    "enum": ["random", "timestamp", "id", "original"]
                },
            },
            "required": ["token"],
        },
        "sqlite": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "default": "disarchive"},
            },
            "required": ["name"],
        },
        "mysql": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "default": "disarchive"},
                "host": {"type": "string", "default": "localhost"},
                "port": {"type": "integer", "default": 3306},
                "user": {"type": "string"},
                "password": {"type": "string"},
            },
            "required": [
                "name",
                "host",
                "port",
                "user",
                "password",
            ],
        },
        "postgresql": {
            "type": "object",
            "properties": {
                "name": {"type": "string", "default": "disarchive"},
                "host": {"type": "string", "default": "localhost"},
                "port": {"type": "integer", "default": 5432},
                "user": {"type": "string"},
                "password": {"type": "string"},
            },
            "required": [
                "name",
                "host",
                "port",
                "user",
                "password",
            ],
        },
    },
    "required": ["general"],
}


# Load config file or alert user if not found
def load_config():
    # create images directory if it doesn't exist
    if not os.path.exists("images"):
        os.makedirs("images")

    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:
        sys.exit(
            LOG.critical(
                "config.yaml file not found. Please use the template to"
                " configure the bot."
            )
        )


# Validate the config file against the schema
def validate_config(file_contents):
    global TOKEN, NAMING_SCHEME, BOT_COLOR, SQLITE_NAME, 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})$"

    # Check if the bot_color is a valid hex color
    if "bot_color" in config["general"]:
        if not bool(
            re.match(hex_pattern_one, config["general"]["bot_color"])
        ) and not bool(
            re.match(hex_pattern_two, config["general"]["bot_color"])
        ):
            LOG.warn(
                "bot_color is not a valid hex color... defaulting to #26dfc9"
            )
        else:
            BOT_COLOR = discord.Color(
                int((config["general"]["bot_color"]).replace("#", ""), 16)
            )

    # Naming scheme
    if "naming_scheme" in config["general"]:
        NAMING_SCHEME = config["general"]["naming_scheme"]
    else:
        LOG.info("No naming scheme specified... defaulting to random")
        NAMING_SCHEME = "random"

    # Assign database variables
    if "sqlite" in config:
        DB_ENGINE = "sqlite"
        if "name" in config["sqlite"]:
            SQLITE_NAME = config["sqlite"]["name"]

    elif "mysql" in config:
        DB_ENGINE = "mysql"
        DB_NAME = config["mysql"]["name"]
        DB_HOST = config["mysql"]["host"]
        DB_PORT = config["mysql"]["port"]
        DB_USER = config["mysql"]["user"]
        DB_PASSWORD = config["mysql"]["password"]

    elif "postgresql" in config:
        DB_ENGINE = "postgresql"
        DB_NAME = config["postgresql"]["name"]
        DB_HOST = config["postgresql"]["host"]
        DB_PORT = config["postgresql"]["port"]
        DB_USER = config["postgresql"]["user"]
        DB_PASSWORD = config["postgresql"]["password"]

    else:
        LOG.warn("No database engine specified. Defaulting to SQLite.")
        DB_ENGINE = "sqlite"

    TOKEN = config["general"]["token"]