import pymysql
import pandas as pd

DB_HOST = 'qa.db.qaorch.com'
DB_USER = ''
DB_PASS = ''
DB_NAME = 'mysql'
DB_SCHEMA = 'art_relations'


def show_grants_for_sha_users(connection):
    try:
        with connection.cursor() as cursor:
            # Query to find users with SHA passwords
            sha_users_sql = "SELECT user, host, authentication_string FROM mysql.user"
            cursor.execute(sha_users_sql)
            sha_users = cursor.fetchall()

            # Save the results to a CSV file
            df = pd.DataFrame(sha_users)
            df.to_csv('sha_users.csv', index=False)

            # Filter users with SHA passwords
            sha_users_filtered = [user for user in sha_users if
                                  user['authentication_string'] and
                                  user[
                                      'authentication_string'].startswith(
                                      'SHA')]

            for user in sha_users_filtered:
                user_name = user['user']
                host = user['host']
                show_grants_sql = f"SHOW GRANTS FOR '{user_name}'@'{host}'"
                cursor.execute(show_grants_sql)
                grants = cursor.fetchall()
                print(f"Grants for {user_name}@{host}:")
                for grant in grants:
                    print(grant[f'Grants for {user_name}@{host}'])
                print(
                    "--------------------------------------------------")
    except Exception as e:
        print(f"An error occurred: {e}")


def generate_sql_from_csv(connection):
    # Read the CSV file
    df = pd.read_csv('sha_users.csv')

    # Generate SQL commands
    create_user_commands = []
    show_grants_results = []
    for index, row in df.iterrows():
        user = row['user']
        host = row['host']
        auth_string = row['authentication_string']

        # Check if user or host is NaN
        if pd.isna(user) or pd.isna(host):
            continue

        create_user_sql = f"CREATE USER IF NOT EXISTS '{user}'@'{host}' IDENTIFIED BY '{auth_string}';"
        create_user_commands.append(create_user_sql)

        # Execute SHOW GRANTS and append the results
        show_grants_sql = f"SHOW GRANTS FOR '{user}'@'{host}';"
        with connection.cursor() as cursor:
            cursor.execute(show_grants_sql)
            grants = cursor.fetchall()
            for grant in grants:
                show_grants_results.append(
                    grant[f'Grants for {user}@{host}'])

    # Save the CREATE USER commands to a .sql file
    with open('create_sha_users.sql', 'w') as file:
        for command in create_user_commands:
            file.write(command + '\n')

    # Save the SHOW GRANTS results to a different .sql file
    with open('show_grants_sha_users.sql', 'w') as file:
        for result in show_grants_results:
            file.write(result + '\n')


# Connect to the database
connection = pymysql.connect(host=DB_HOST, user=DB_USER,
                             password=DB_PASS, db=DB_NAME,
                             charset='utf8mb4',
                             cursorclass=pymysql.cursors.DictCursor)

try:
    with connection.cursor() as cursor:
        '''Get user permissions'''
        user_audit_sql = 'select host, user, db, insert_priv, ' \
                         'update_priv, ' \
                         'delete_priv, create_tmp_table_priv, alter_priv ' \
                         'from mysql.db;'

        '''Get schema permissions'''
        schema_audit_sql = 'select user, select_priv, insert_priv, ' \
                           'update_priv, delete_priv, create_priv, drop_priv ' \
                           'from mysql.db where db = "art_relations";'

        '''Get table permissions'''
        table_audit_sql = 'select user, table_name, table_priv ' \
                          'from mysql.tables_priv where db = "art_relations";'

        for sql in [user_audit_sql, schema_audit_sql, table_audit_sql]:
            if sql == user_audit_sql:
                cursor.execute(user_audit_sql)
                result = cursor.fetchall()
                df = pd.DataFrame(result)
                df.to_csv('user_audit_sql.csv', index=False)
            elif sql == schema_audit_sql:
                cursor.execute(schema_audit_sql)
                result = cursor.fetchall()
                df = pd.DataFrame(result)
                df.to_csv('schema_audit_sql.csv', index=False)
            elif sql == table_audit_sql:
                cursor.execute(table_audit_sql)
                result = cursor.fetchall()
                df = pd.DataFrame(result)
                df.to_csv('table_audit_sql.csv', index=False)
            else:
                print('No results found')

    # Show grants for users with SHA passwords
    show_grants_for_sha_users(connection)

    # Generate SQL commands from CSV and save to .sql files
    generate_sql_from_csv(connection)

finally:
    connection.close()