import yaml
from pysnmp.hlapi import *
from rich.console import Console
from rich.markup import escape # Adicionada importação para corrigir o erro de markup

# Inicializa o console do Rich para uma saída colorida e bonita
console = Console()

# --- CONFIGURAÇÃO ---
# Lista de communities para testar em caso de falha da principal
# CORREÇÃO: Corrigido erro de digitação de 'COMMUNenties' para 'COMMUNITIES'
FALLBACK_COMMUNITIES = ["mgr4clientetche", "mgr4rnp!"]
# OID padrão para verificar se a community está correta (sysDescr)
TEST_OID = '1.3.6.1.2.1.1.1.0'
# Timeout da requisição em segundos - AUMENTADO PARA SER MAIS TOLERANTE
SNMP_TIMEOUT = 5

def load_hosts_config(filename='hosts.yml'):
    """Carrega a configuração dos hosts a partir de um arquivo YAML."""
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            return yaml.safe_load(file)
    except FileNotFoundError:
        console.print(f"[bold red]Erro: O arquivo de configuração '{filename}' não foi encontrado.[/bold red]")
        return None
    except yaml.YAMLError as e:
        console.print(f"[bold red]Erro de formatação no arquivo YAML '{filename}': {e}[/bold red]")
        return None

def test_single_community(host_ip, community):
    """
    Tenta executar uma única consulta SNMP a um host com uma community específica.
    Retorna True se for bem-sucedido, False caso contrário.
    """
    errorIndication, errorStatus, errorIndex, varBinds = next(
        getCmd(SnmpEngine(),
               CommunityData(community, mpModel=1),
               UdpTransportTarget((host_ip, 161), timeout=SNMP_TIMEOUT, retries=0),
               ContextData(),
               ObjectType(ObjectIdentity(TEST_OID))),
        (None, None, None, None) # Valor padrão em caso de timeout
    )

    # Verifica se houve erro na comunicação ou na consulta
    if errorIndication or errorStatus:
        return False
    # Verifica se a resposta não está vazia (caso de timeout)
    if not varBinds:
        return False
        
    return True

def main():
    """Função principal do script."""
    console.print("[bold cyan]Iniciando verificação das communities SNMP...[/bold cyan]")
    
    hosts = load_hosts_config()
    if not hosts:
        return

    # Itera sobre cada host configurado no arquivo .yml
    for host_key, config in hosts.items():
        host_ip = config.get('ip')
        primary_community = config.get('community')
        display_name = config.get('displayName', host_key)

        if not host_ip:
            console.print(f"-> [yellow]Aviso:[/yellow] Host '{display_name}' não possui um IP configurado. Pulando.")
            continue

        console.print(f"\n[bold]Testando:[/bold] {display_name} ({host_ip})")

        # Cria a lista de communities a serem testadas, com a primária primeiro
        communities_to_test = []
        if primary_community:
            communities_to_test.append(primary_community)
        communities_to_test.extend(FALLBACK_COMMUNITIES)
        
        found_working_community = False
        # Testa cada community da lista
        for community in communities_to_test:
            console.print(f"  - Tentando community: '{community}'...", end="")
            if test_single_community(host_ip, community):
                console.print(f" [bold green]SUCESSO![/bold green]")
                found_working_community = True
                break # Para o teste para este host, pois já encontramos uma community válida
            else:
                console.print(f" [red]Falhou.[/red]")
        
        if not found_working_community:
            # CORREÇÃO: A linha que causa o erro agora está dentro de um bloco
            # try/except. Se a biblioteca de formatação falhar, o script
            # imprimirá uma mensagem simples e continuará para o próximo host
            # sem travar.
            try:
                console.print(f"[bold red]ATENÇÃO:[/bold] Nenhuma community funcionou para", escape(display_name))
            except Exception:
                # Fallback para um print simples para garantir que o script não trave
                print(f"\nATENÇÃO: Nenhuma community funcionou para {display_name}")


if __name__ == '__main__':
    main()