1#@!#!123s

��JFIF�� 1#@!#!123s
: / opt / sharedrads / cms_tools /
Filename : db.py
back
#! /opt/imh-python/bin/python3
""" Database functions for working with CMS. """
# Author: Daniel K

import os
import re
import logging
import pymysql
import glob


from rads import prompt_y_n
from cms_tools.helpers import (
    common_get_string,
    make_valid_db_name,
    import_db,
    dump_db,
    db_exists,
    db_user_exists,
    create_db,
    create_db_user,
    change_db_pass,
    associate_db_user,
    get_mysql_err,
)

from cms_tools.cms import CMSStatus, CMSError


LOGGER = logging.getLogger(__name__)


def import_cms_db(the_cms, database_file):
    '''
    Import database, using known credentials.
    Also, make backup if necessary.
    '''

    LOGGER.debug(
        "Attempting to import %s into %s", database_file, the_cms.db_name
    )

    database = the_cms.db_name

    if database not in the_cms.modified_dbs:
        dump_file = dump_db(
            the_cms.db_user,
            the_cms.db_pass,
            the_cms.db_name,
            the_cms.directory_root,
            'cms_tools_backup',
        )
        if not dump_file:
            the_cms.set_status(
                CMSStatus.critical, "Unable to backup %s" % the_cms.db_name
            )
            raise CMSError("Unable to backup %s" % the_cms.db_name)

        LOGGER.info("Modifying database %s", database)
        the_cms.modified_dbs[database] = dump_file

    return import_db(
        the_cms.db_user, the_cms.db_pass, the_cms.db_name, database_file
    )


def test_db_connection(the_cms):
    '''
    Test the database connection and return errors
    '''
    LOGGER.debug("Testing db connection")
    try:
        with pymysql.connect(
            host=the_cms.db_host,
            user=the_cms.db_user,
            password=the_cms.db_pass,
            database=the_cms.db_name,
        ) as conn:
            with conn.cursor() as cursor:
                if cursor.execute("SHOW TABLES") == 0:
                    return None  # No tables, but connected
    except pymysql.Error as err:
        LOGGER.debug("Connection error")
        return err
    return None


def simple_query(the_cms, field, table, search_field='', search_pattern=''):
    '''
    Search database returning specific field with optional search parameters
    '''

    if the_cms.status < CMSStatus.db_has_matching_tables:
        LOGGER.error(
            "Database %s has not yet been confrmed working, "
            "but query attempted",
            the_cms.db_name,
        )
        return None

    if search_pattern == '':
        if search_field != '':
            LOGGER.warning("Search field given, but no pattern given")
            search_field = ''
    # MySQL identifiers can't be escaped by execute() like literals can
    prefix = the_cms.db_pref.replace('`', '``')
    escaped_table = f"`{prefix}{table.replace('`', '``')}`"
    escaped_field = f"`{field.replace('`', '``')}`"
    if search_pattern == '':
        query = f"SELECT {escaped_field} FROM {escaped_table};"
        args = None
    else:
        escaped_search = f"`{search_field.replace('`', '``')}`"
        args = (search_pattern,)
        query = (
            f"SELECT {escaped_field} FROM {escaped_table} "
            f"WHERE {escaped_search} LIKE %s"
        )
    try:
        with pymysql.connect(
            host=the_cms.db_host,
            user=the_cms.db_user,
            password=the_cms.db_pass,
            database=the_cms.db_name,
        ) as conn:
            with conn.cursor() as cursor:
                result = cursor.execute(query, args)
                if result < 1:  # No tables, but connected
                    return None
                return cursor.fetchall()[0]
    except pymysql.Error as err:
        LOGGER.error(err)
        return None


def check_db_auth(the_cms):
    '''
    Check whether the database user and password is correct
    '''

    # First, see whether the name uses a valid format
    if not re.match(
        "%s_[a-z0-9]{1,%d}" % (the_cms.dbprefix, 15 - len(the_cms.dbprefix)),
        the_cms.db_user,
    ):
        LOGGER.info("Database username '%s' is not correct", the_cms.db_user)

        new_name = make_valid_db_name(
            the_cms.cpuser, the_cms.dbprefix, the_cms.db_user, name_type="user"
        )
        if None is new_name:

            # If we did not get a new name, allow the user to make one
            new_name = common_get_string(
                "What new name would you like? ",
                "%s_[a-z0-9]{1,%d}"
                % (the_cms.dbprefix, 15 - len(the_cms.dbprefix)),
            )
            if None is not new_name:
                if not the_cms.set_variable('db_user', new_name):
                    return False
                the_cms.db_user = the_cms.get_variable('db_user')
                LOGGER.info("Username set to %s", the_cms.db_user)
            else:
                LOGGER.error("Username '%s' not reset!", the_cms.db_user)
                return False

        else:

            # We got a new name. Prompt to set it
            if the_cms.ilevel < 1 or prompt_y_n("Set name to %s? " % new_name):
                if not the_cms.set_variable('db_user', new_name):
                    return False
                the_cms.db_user = the_cms.get_variable('db_user')
                LOGGER.info("Username set to %s", the_cms.db_user)
            else:
                new_name = common_get_string(
                    "Use what database user name: ", 'database'
                )
                if not the_cms.set_variable('db_user', new_name):
                    return False
                the_cms.db_user = the_cms.get_variable('db_user')
                LOGGER.info("Username set to %s", the_cms.db_user)

    # Username is a valid format

    # Does it exist and work?

    # We can just check whether it exists, and if not, create it
    if not db_user_exists(the_cms.cpuser, the_cms.db_user):
        if the_cms.ilevel < 1 or prompt_y_n(
            "Database user '%s' does not exist. Create it?" % the_cms.db_user
        ):
            create_db_user(the_cms.cpuser, the_cms.db_user, the_cms.db_pass)

            # Check just in case it's not really added
            if not db_user_exists(the_cms.cpuser, the_cms.db_user):
                the_cms.set_status(
                    CMSStatus.error,
                    "Failed to create database user '%s'" % the_cms.db_user,
                )
                return False
        else:
            the_cms.set_status(
                CMSStatus.error, "Could not create %s" % the_cms.db_user
            )
            return False

    # So, the db user exists. Does the pw match?
    result = test_db_connection(the_cms)
    if result is None:
        LOGGER.debug("Authorization fixed")
        return True
    if 1045 == result[0]:
        if the_cms.ilevel < 1 or prompt_y_n(
            "Password for user '%s' doesn't match. Reset it?" % the_cms.db_user
        ):
            if not change_db_pass(
                the_cms.cpuser, the_cms.db_user, the_cms.db_pass
            ):
                LOGGER.error("Could reset password for %s.", the_cms.db_user)
            return True
        LOGGER.error("Could not fix password for %s.", the_cms.db_user)
        return False
    if 1044 == result[0]:
        # The user isn't associated, but this confirms auth worked
        return True
    # Some other error, so we'll assume this is not the issue
    (errno, sterror) = result
    LOGGER.info(
        "Database connection failing. "
        "Can't check username. "
        "Receiving error:\n(%d): %s",
        errno,
        sterror,
    )
    return True


# End check_db_auth


def check_db_access(the_cms):
    '''
    Check whether the database exists and the user has privileges
    '''

    # First, see whether the name uses a valid format
    if not re.match(
        "%s_[a-z0-9]{1,%d}" % (the_cms.dbprefix, 15 - len(the_cms.dbprefix)),
        the_cms.db_name,
    ):
        LOGGER.info("Database name '%s' is not correct", the_cms.db_name)

        new_name = make_valid_db_name(
            the_cms.cpuser, the_cms.dbprefix, the_cms.db_name
        )
        if None is new_name:

            # If we did not get a new name, allow the user to make one
            new_name = common_get_string(
                "What new database name would you like? ",
                "%s_[a-z0-9]{1,%d}"
                % (the_cms.dbprefix, 15 - len(the_cms.dbprefix)),
            )
            if None is not new_name:
                if not the_cms.set_variable('db_name', new_name):
                    return False
                the_cms.db_name = the_cms.get_variable('db_name')
                LOGGER.info("Database name set to %s", the_cms.db_name)
            else:
                the_cms.set_status(CMSStatus.error, "Database name not correct")
                return False

        else:

            # We got a new name. Prompt to set it
            if the_cms.ilevel < 1 or prompt_y_n("Set name to %s?" % new_name):
                if not the_cms.set_variable('db_name', new_name):
                    return False
                the_cms.db_name = the_cms.get_variable('db_name')
                LOGGER.info("Database name set to %s", the_cms.db_name)
            else:
                new_name = common_get_string(
                    "Use what database name: ", 'database'
                )
                if not the_cms.set_variable('db_name', new_name):
                    return False
                the_cms.db_name = the_cms.get_variable('db_name')
                LOGGER.info("Database name set to %s", the_cms.db_name)

    # Database name is a valid format

    if not db_exists(the_cms.cpuser, the_cms.db_name):
        if the_cms.ilevel < 1 or prompt_y_n(
            "Database '%s' does not exist. Create it?" % the_cms.db_name
        ):
            create_db(the_cms.cpuser, the_cms.db_name)
            if not db_exists(the_cms.cpuser, the_cms.db_name):
                the_cms.set_status(
                    CMSStatus.error,
                    "Failed to create database '%s'" % the_cms.db_name,
                )
                return False
        else:
            the_cms.set_status(CMSStatus.error, "Database could not be created")
            return False

    # Did adding the database fix the problem?
    result = test_db_connection(the_cms)
    if result is None:
        # Yes, that did it
        LOGGER.debug("Database connection fixed.")
        return True
    errno, sterror = get_mysql_err(result)
    if errno not in (1044, 1049):
        # Not certain, but not the same error, so pretend that it did.
        LOGGER.error(
            "Still could not connect to '%s'. New error: %d: %s",
            the_cms.db_name,
            errno,
            sterror,
        )
        return True

    LOGGER.debug(
        "The database exists, but the user cannot access the database."
    )

    # If we've made it here, we can assign the user
    if the_cms.ilevel < 1 or prompt_y_n(
        "Associate database user '%s' with database '%s'?"
        % (the_cms.db_user, the_cms.db_name)
    ):
        associate_db_user(the_cms.cpuser, the_cms.db_name, the_cms.db_user)
    else:
        the_cms.set_status(
            CMSStatus.error, "Could not associate database user."
        )
        LOGGER.warning("Could not associate database user.")
        return False

    return True


# End check_db_access


def check_db_error(the_cms):
    '''
    Check for database connection errors.
    Return None if no error or number if there was an errror
    '''

    # Make sure everything was set up first
    if the_cms.status < CMSStatus.db_is_set:
        LOGGER.warning(
            "Database credentials haven't been set. Last status: %s",
            the_cms.reason,
        )
        return -1

    # Make sure that we're checking the local db first
    if 'localhost' != the_cms.db_host:
        LOGGER.info("Databse host is set to '%s'.", the_cms.db_host)
        if the_cms.ilevel < 1 or prompt_y_n("Set database host to localhost?"):
            if not the_cms.set_variable('db_host', "localhost"):
                return -1
            the_cms.db_host = the_cms.get_variable('db_host')
            LOGGER.debug("Database host has been fixed")

    result = test_db_connection(the_cms)
    if result is None:
        LOGGER.debug("Database connection working")
        return None
    return get_mysql_err(result)[0]


# End check_db_error


def fix_db_error(the_cms, error_number):
    '''
    Check whether the database connection is working
    '''

    if error_number is None:
        return True

    if error_number == -1:
        return False

    LOGGER.info("There was a database error for %s", the_cms.db_name)

    if error_number == 1045:
        LOGGER.info("The username or password is incorrect")
        return check_db_auth(the_cms)
    if error_number in (1044, 1049):
        LOGGER.info("The user cannot access the database")
        return check_db_access(the_cms)
    if error_number == 2006:
        LOGGER.info(
            "MySQL server has gone away. May need to be researched manually"
        )
        return False
    # Unknown error
    LOGGER.error("Unknown error.")
    LOGGER.error(error_number)
    return False


# End fix_db_error


def check_db(the_cms):
    '''
    Check whether the database connection is working
    '''

    # Make sure everything was set up first
    if the_cms.status < CMSStatus.db_is_set:
        LOGGER.warning(
            "Database credentials haven't been set. Last status: %s",
            the_cms.reason,
        )
        return False

    # Make sure that we're checking the local db first
    if 'localhost' != the_cms.db_host:
        LOGGER.info("Databse host is set to '%s'.", the_cms.db_host)
        if the_cms.ilevel < 1 or prompt_y_n("Set database host to localhost?"):
            if not the_cms.set_variable('db_host', "localhost"):
                return False
            the_cms.db_host = the_cms.get_variable('db_host')
            LOGGER.debug("Database host has been fixed")

    db_error = check_db_error(the_cms)
    count = 0
    while None is not db_error:
        if not fix_db_error(the_cms, db_error):
            LOGGER.info("Could not resolve database error %s", db_error)
            return False

        count += 1
        if count > 10:
            LOGGER.error("Too many database errors. Giving up.")
            return False

        db_error = check_db_error(the_cms)

    the_cms.set_status(
        CMSStatus.db_is_connecting, "Database confirmed connected"
    )
    return True


# End check_db


def check_db_data(the_cms):
    '''
    Check database to ensure that is not empty, and that it has tables
    matching the prefix. If not, attempt to import.
    '''

    if the_cms.status < CMSStatus.db_is_connecting:
        if not check_db(the_cms):
            LOGGER.warning(
                "Database has not been confirmed to connect. "
                "Cannot check database data."
            )
            return False

    try:
        with pymysql.connect(
            host=the_cms.db_host,
            user=the_cms.db_user,
            password=the_cms.db_pass,
            database=the_cms.db_name,
        ) as conn:
            with conn.cursor() as cursor:
                if cursor.execute("SHOW TABLES") == 0:
                    LOGGER.info("No tables in '%s'", the_cms.db_name)
                    return fix_empty_db(the_cms)
                the_cms.set_status(
                    CMSStatus.db_has_tables, "Database has tables"
                )
                count = cursor.execute(
                    "SHOW TABLES LIKE %s%%", (the_cms.db_pref,)
                )
                if count == 0:
                    LOGGER.info(
                        "Database '%s' has tables, "
                        "but none matching the '%s' prefix.",
                        the_cms.db_name,
                        the_cms.db_pref,
                    )
                    return fix_empty_db(the_cms)
    except pymysql.Error as err:
        raise CMSError(f"Database query error: {err}") from err
    the_cms.set_status(
        CMSStatus.db_has_matching_tables,
        "Database '{}' has tables matching the '{}' prefix.".format(
            the_cms.db_name, the_cms.db_pref
        ),
    )
    return True


def fix_db_names(the_cms):
    '''
    Set database name and database user names to proper names
    which fit with cPanel
    '''

    # Set new names
    new_db_name = make_valid_db_name(
        the_cms.cpuser, the_cms.dbprefix, the_cms.db_name
    )
    new_db_user = make_valid_db_name(
        the_cms.dbuser, the_cms.dbprefix, the_cms.db_user, name_type="user"
    )

    the_cms.set_variable('db_name', new_db_name)
    the_cms.set_variable('db_user', new_db_user)


# End fix_db_names()


def fix_empty_db(the_cms):
    '''
    Attempt to find and import database
    '''

    if the_cms.status < CMSStatus.db_is_connecting:
        if not check_db(the_cms):
            LOGGER.warning(
                "Database not confirmed connecting. Cannot import data"
            )
            return False

    found_files = []
    for dbname in (the_cms.orig_db_name, the_cms.db_name):
        for the_directory in the_cms.db_file_search_path:
            for db_file in glob.glob(
                os.path.join(the_directory, "*%s*.sql" % dbname)
            ):

                if db_file in found_files:
                    continue

                found_files.append(db_file)
                if the_cms.ilevel < 1 or prompt_y_n(
                    f"Import {db_file} in to {the_cms.db_name}?"
                ):
                    return import_cms_db(the_cms, db_file)

    if not the_cms.ilevel < 1:
        db_file = common_get_string(
            "Please specify a database "
            "to import into %s: " % the_cms.db_name,
            default=None,
        )
        while db_file is not None:
            if os.path.isfile(db_file):
                return import_cms_db(the_cms, db_file)
            db_file = common_get_string(
                "File does no exist. Please specify a "
                f"database to import into {the_cms.db_name}: ",
                default=None,
            )

    the_cms.set_status(
        CMSStatus.warning,
        f"Could not find a database export for {the_cms.db_name}",
    )
    return False