import psycopg2
import re
import yaml

def generate_password(length=8):
    # Определяем возможные символы для разных категорий
    lower_case = string.ascii_lowercase
    upper_case = string.ascii_uppercase
    digits = string.digits
    special_characters = string.punctuation

    # Гарантируем, что в пароле будет хотя бы один символ каждого типа
    password = [
        secrets.choice(lower_case),      # строчная буква
        secrets.choice(upper_case),      # заглавная буква
        secrets.choice(digits),          # цифра
        secrets.choice(special_characters) # спецсимвол
    ]

    # Заполняем оставшиеся символы случайными из всех категорий
    all_characters = lower_case + upper_case + digits + special_characters
    password += [secrets.choice(all_characters) for _ in range(length - 4)]

    # Перемешиваем полученные символы
    secrets.SystemRandom().shuffle(password)

    # Собираем строку из символов и возвращаем
    return ''.join(password)

# Функция для чтения конфигурации из YAML файла
def get_configurations(properties_file):
    with open(properties_file, 'r') as file:
        config = yaml.safe_load(file)  # Чтение YAML файла

        # Извлекаем данные из конфигурации
        db_url = config.get('spring', {}).get('datasource', {}).get('url', '')
        db_user = config.get('spring', {}).get('datasource', {}).get('username', '')
        current_password = config.get('spring', {}).get('datasource', {}).get('password', '')

        # Разбираем URL для получения host, port, и dbname
        match = re.match(r'jdbc:postgresql://([^:]+):(\d+)/([^?]+)', db_url)
        if match:
            db_host = match.group(1)
            db_port = match.group(2)
            db_name = match.group(3)
        else:
            raise ValueError("Database URL is invalid in the properties file.")

        return db_host, db_port, db_name, db_user, current_password

# Функция для обновления пароля в файле конфигурации
def update_password_in_file(properties_file, current_password, new_password):
    with open(properties_file, 'r') as file:
        content = file.read()

    # Заменяем старый пароль на новый
    updated_content = re.sub(rf'spring.datasource.password={current_password}', f'spring.datasource.password={new_password}', content)

    with open(properties_file, 'w') as file:
        file.write(updated_content)

# Основная функция для смены пароля в PostgreSQL
def change_postgresql_password(db_host, db_port, db_name, db_user, current_password, new_password):
    try:
        # Подключаемся к PostgreSQL с использованием старого пароля
        conn = psycopg2.connect(
            dbname=db_name,
            user=db_user,
            password=current_password,
            host=db_host,
            port=db_port
        )
        conn.autocommit = True

        # Создаем курсор для выполнения SQL-запроса
        cursor = conn.cursor()

        # Запрос для изменения пароля пользователя
        cursor.execute(f"ALTER USER {db_user} WITH PASSWORD '{new_password}';")

        cursor.close()
        conn.close()

        print(f"Password for user {db_user} successfully changed.")
    except Exception as e:
        print(f"Error changing password in PostgreSQL: {e}")
        return False
    return True

# Главная логика скрипта
def main():
    # Путь к файлу конфигурации Spring Boot
    properties_file = "application.yml"

    # Получаем данные из файла конфигурации
    db_host, db_port, db_name, db_user, current_password = get_configurations(properties_file)

    # Новый пароль
    new_password = generate_password()  # Укажите новый пароль здесь

    # Меняем пароль в PostgreSQL
    if change_postgresql_password(db_host, db_port, db_name, db_user, current_password, new_password):
        # Обновляем пароль в файле конфигурации
        update_password_in_file(properties_file, current_password, new_password)
        print(f"Password updated successfully in {properties_file}.")
    else:
        print("Password change failed.")

if __name__ == "__main__":
    main()