bgp-server/scripts/load_prefixes.py

#!/usr/bin/env python3
"""
Скрипт для загрузки префиксов из текстовых файлов в ExaBGP
"""

import sys
import os
import time
import json
import ipaddress
from pathlib import Path
from typing import List, Set

# Настройки
DATA_DIR = "/app/data"
LOG_FILE = "/app/logs/prefixes.log"
PREFIX_FILES = ["prefixes.txt", "additional_prefixes.txt"]

def log_message(message: str):
    """Запись сообщения в лог"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"[{timestamp}] {message}"
    
    # Вывод в stdout для ExaBGP
    print(log_entry, flush=True)
    
    # Запись в файл (с обработкой ошибок)
    try:
        with open(LOG_FILE, "a") as f:
            f.write(log_entry + "\n")
    except (PermissionError, OSError) as e:
        # Если не можем записать в файл, просто выводим в stderr
        print(f"Warning: Could not write to log file: {e}", file=sys.stderr, flush=True)

def validate_prefix(prefix: str) -> bool:
    """Валидация IP префикса"""
    try:
        ipaddress.ip_network(prefix, strict=False)
        return True
    except ValueError:
        return False

def load_prefixes_from_file(file_path: Path) -> Set[str]:
    """Загрузка префиксов из файла"""
    prefixes = set()
    
    if not file_path.exists():
        log_message(f"File not found: {file_path}")
        return prefixes
    
    try:
        with open(file_path, 'r') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                
                # Пропуск пустых строк и комментариев
                if not line or line.startswith('#'):
                    continue
                
                # Валидация префикса
                if validate_prefix(line):
                    prefixes.add(line)
                else:
                    log_message(f"Invalid prefix in {file_path}:{line_num}: {line}")
    
    except Exception as e:
        log_message(f"Error reading {file_path}: {e}")
    
    return prefixes

def load_all_prefixes() -> Set[str]:
    """Загрузка всех префиксов из всех файлов"""
    all_prefixes = set()
    
    for filename in PREFIX_FILES:
        file_path = Path(DATA_DIR) / filename
        prefixes = load_prefixes_from_file(file_path)
        all_prefixes.update(prefixes)
        log_message(f"Loaded {len(prefixes)} prefixes from {filename}")
    
    return all_prefixes

def announce_prefixes(prefixes: Set[str]):
    """Анонсирование префиксов через ExaBGP API"""
    # Получение настроек из переменных окружения
    neighbor_ip = os.environ.get('BGP_NEIGHBOR_IP', '192.168.1.1')
    local_as = int(os.environ.get('BGP_LOCAL_AS', '65000'))
    router_id = os.environ.get('BGP_ROUTER_ID', '192.168.1.100')
    
    for prefix in sorted(prefixes):
        # Формат команды для ExaBGP
        command = {
            "command": "announce route",
            "neighbor": {
                "ip": neighbor_ip,
                "description": "RouterOS peer"
            },
            "attribute": {
                "origin": "igp",
                "as-path": [local_as],
                "next-hop": router_id
            },
            "nlri": prefix
        }
        
        # Отправка команды в ExaBGP
        print(json.dumps(command), flush=True)
        log_message(f"Announced prefix: {prefix}")

def withdraw_prefixes(prefixes: Set[str]):
    """Отзыв префиксов через ExaBGP API"""
    # Получение настроек из переменных окружения
    neighbor_ip = os.environ.get('BGP_NEIGHBOR_IP', '192.168.1.1')
    
    for prefix in sorted(prefixes):
        command = {
            "command": "withdraw route",
            "neighbor": {
                "ip": neighbor_ip,
                "description": "RouterOS peer"
            },
            "nlri": prefix
        }
        
        print(json.dumps(command), flush=True)
        log_message(f"Withdrew prefix: {prefix}")

def main():
    """Основная функция"""
    # Создание директории для логов в начале
    try:
        os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
    except (PermissionError, OSError):
        pass  # Игнорируем ошибки создания директории
    
    # Проверка тестового режима
    if len(sys.argv) > 1 and sys.argv[1] == '--test':
        print("Testing prefix loader...")
        prefixes = load_all_prefixes()
        print(f"Found {len(prefixes)} prefixes: {list(prefixes)}")
        return
    
    log_message("Starting prefix loader...")
    
    # Загрузка префиксов
    prefixes = load_all_prefixes()
    log_message(f"Total prefixes loaded: {len(prefixes)}")
    
    if not prefixes:
        log_message("No valid prefixes found. Exiting.")
        return
    
    # Анонсирование префиксов
    announce_prefixes(prefixes)
    log_message("All prefixes announced successfully")
    
    # Ожидание команд от ExaBGP
    try:
        while True:
            line = sys.stdin.readline()
            if not line:
                break
            
            # Обработка команд от ExaBGP
            try:
                data = json.loads(line.strip())
                log_message(f"Received command: {data}")
            except json.JSONDecodeError:
                log_message(f"Invalid JSON received: {line.strip()}")
    
    except KeyboardInterrupt:
        log_message("Received interrupt signal")
    except Exception as e:
        log_message(f"Error in main loop: {e}")
    finally:
        # Отзыв всех префиксов при завершении
        log_message("Withdrawing all prefixes...")
        withdraw_prefixes(prefixes)
        log_message("Prefix loader stopped")

if __name__ == "__main__":
    main()