#!/usr/bin/env python3
import sys
import os
import re
import subprocess
import ipaddress

def main():
    # 1. Проверка прав root
    if os.geteuid() != 0:
        print("Ошибка: Скрипт должен запускаться с правами root (sudo).")
        sys.exit(1)

    # 2. Проверка аргументов
    if len(sys.argv) != 3:
        print("Использование: setourip -ник_админа IP_админа")
        print("Пример: setourip -odina 188.187.103.53")
        sys.exit(1)

    raw_nick = sys.argv[1]
    new_ip = sys.argv[2]

    # Удаляем ведущий тире, если он был передан (-odina -> odina)
    nick = raw_nick.lstrip('-')

    # Валидация корректности IP-адреса
    try:
        ipaddress.ip_address(new_ip)
    except ValueError:
        print(f"Ошибка: '{new_ip}' не является валидным IP-адресом.")
        sys.exit(1)

    file_txt = "/etc/haproxy/whitelists/whitelist-our-ips.txt"
    file_yaml = "/etc/crowdsec/parsers/s02-enrich/whitelist-our-ips.yaml"

    updated = False

    # --- 3. Обновление TXT (HAProxy) ---
    if os.path.exists(file_txt):
        with open(file_txt, 'r', encoding='utf-8') as f:
            content_txt = f.read()

        # Ищем паттерн: #ник, перенос строки и следом IP
        pattern_txt = re.compile(rf'(#{re.escape(nick)}\s*\n)\s*[0-9a-fA-F.:]+', re.MULTILINE)
        
        if pattern_txt.search(content_txt):
            new_content_txt = pattern_txt.sub(rf'\g<1>{new_ip}', content_txt)
            with open(file_txt, 'w', encoding='utf-8') as f:
                f.write(new_content_txt)
            print(f"[OK] {file_txt} обновлен для '{nick}'.")
            updated = True
        else:
            print(f"[WARNING] Ник #{nick} не найден в {file_txt}.")
    else:
        print(f"[ERROR] Файл {file_txt} не найден.")

    # --- 4. Обновление YAML (CrowdSec) ---
    if os.path.exists(file_yaml):
        with open(file_yaml, 'r', encoding='utf-8') as f:
            content_yaml = f.read()

        # Ищем паттерн: - "IP" #ник (с учетом пробелов)
        pattern_yaml = re.compile(rf'-\s*"[^"]*"\s*#{re.escape(nick)}')
        
        if pattern_yaml.search(content_yaml):
            new_content_yaml = pattern_yaml.sub(f'- "{new_ip}" #{nick}', content_yaml)
            with open(file_yaml, 'w', encoding='utf-8') as f:
                f.write(new_content_yaml)
            print(f"[OK] {file_yaml} обновлен для '{nick}'.")
            updated = True
        else:
            print(f"[WARNING] Ник #{nick} не найден в {file_yaml}.")
    else:
        print(f"[ERROR] Файл {file_yaml} не найден.")

    # --- 5. Проверка конфигурации HAProxy ---
    if updated:
        print("\nПроверка конфигурации HAProxy...")
        try:
            # Выполняем синтаксическую проверку конфигурации HAProxy
            check_res = subprocess.run(
                ["haproxy", "-c", "-f", "/etc/haproxy/haproxy.cfg"],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            print("[OK] Конфигурация HAProxy корректна.")
        except subprocess.CalledProcessError as e:
            print(f"[CRITICAL ERROR] Ошибка в конфигурации HAProxy:\n{e.stderr}")
            print("[CANCELLED] Перезагрузка сервисов ОТМЕНЕНА из-за ошибки в конфигурации HAProxy!")
            sys.exit(1)

        # --- 6. Перезагрузка сервисов ---
        print("\nПерезагрузка сервисов...")
        
        # Перезагрузка HAProxy
        try:
            subprocess.run(["systemctl", "reload", "haproxy"], check=True)
            print("[OK] HAProxy успешно перезагружен (reload).")
        except subprocess.CalledProcessError as e:
            print(f"[ERROR] Не удалось перезагрузить HAProxy: {e}")

        # Перезагрузка CrowdSec
        try:
            subprocess.run(["systemctl", "reload", "crowdsec"], check=True)
            print("[OK] CrowdSec успешно перезагружен (reload).")
        except subprocess.CalledProcessError as e:
            print(f"[ERROR] Не удалось перезагрузить CrowdSec: {e}")

if __name__ == "__main__":
    main()