#!/usr/bin/env python3
"""
gNMI gRPC Configuration Script
===============================
Enables gRPC dial-in telemetry on all 9 IOS-XR routers so that
Telegraf (or any gNMI collector) can subscribe to streaming
telemetry data.

What this script applies per router:
  - gRPC server on port 57400 with TLS disabled

Uses SSH/CLI (paramiko) instead of NETCONF because IOS-XR 24.3.1
rejects the NETCONF edit-config for gRPC with "Need to enable GRPC first".

Router targets:
  CORE-01 (10.100.0.100)
  CORE-02 (10.100.0.200)
  R9K-01  (10.100.0.1) through R9K-07 (10.100.0.7)
"""

import paramiko
import time
import sys

ROUTERS = [
    ('10.100.0.100', 'CORE-01'),
    ('10.100.0.200', 'CORE-02'),
    ('10.100.0.1',   'R9K-01'),
    ('10.100.0.2',   'R9K-02'),
    ('10.100.0.3',   'R9K-03'),
    ('10.100.0.4',   'R9K-04'),
    ('10.100.0.5',   'R9K-05'),
    ('10.100.0.6',   'R9K-06'),
    ('10.100.0.7',   'R9K-07'),
]

USERNAME = 'webui'
PASSWORD = 'cisco'
GRPC_PORT = 57400

CONFIG_COMMANDS = [
    'configure terminal',
    'grpc',
    f'port {GRPC_PORT}',
    'no-tls',
    'commit',
    'end',
]


def configure_router(mgmt_ip, label):
    """Apply gRPC configuration via SSH CLI."""
    print(f"\n{'─'*60}")
    print(f"  Configuring {label} ({mgmt_ip})")
    print(f"{'─'*60}")
    print(f"  Applying: gRPC port={GRPC_PORT} no-tls")

    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(mgmt_ip, username=USERNAME, password=PASSWORD, timeout=10)

        shell = client.invoke_shell()
        time.sleep(1)
        shell.recv(65535)  # clear banner

        for cmd in CONFIG_COMMANDS:
            shell.send(cmd + '\n')
            time.sleep(1.5)

        output = shell.recv(65535).decode()
        client.close()

        if 'error' in output.lower() or 'fail' in output.lower():
            print(f"  ✗ ERROR on {label}: {output.strip()}")
            return False

        print(f"  ✓ {label} done.")
        return True

    except Exception as e:
        print(f"  ✗ ERROR on {label}: {e}")
        return False


def verify_router(mgmt_ip, label):
    """Verify gRPC configuration via SSH."""
    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(mgmt_ip, username=USERNAME, password=PASSWORD, timeout=10)

        shell = client.invoke_shell()
        time.sleep(1)
        shell.recv(65535)

        shell.send('show running-config grpc\n')
        time.sleep(3)
        output = shell.recv(65535).decode()
        client.close()

        has_port = f'port {GRPC_PORT}' in output
        has_notls = 'no-tls' in output

        p = '✓' if has_port else '✗'
        t = '✓' if has_notls else '✗'
        status = 'OK' if (has_port and has_notls) else 'INCOMPLETE'
        print(f"  {label:8s}  grpc-port={p}  no-tls={t}  [{status}]")
        return has_port and has_notls

    except Exception as e:
        print(f"  {label:8s}  verify error: {e}")
        return False


def main():
    print("gNMI gRPC Configuration Script")
    print("================================")
    print(f"Targets: all {len(ROUTERS)} routers")
    print()

    results = []
    for mgmt_ip, label in ROUTERS:
        ok = configure_router(mgmt_ip, label)
        results.append((mgmt_ip, label, ok))

    # Verification pass
    print(f"\n{'='*60}")
    print("Post-apply verification")
    print('='*60)
    print(f"  {'Router':8s}  {'gRPC Port':9s}  {'No-TLS':6s}  Status")

    all_ok = True
    for mgmt_ip, label, applied_ok in results:
        if applied_ok:
            if not verify_router(mgmt_ip, label):
                all_ok = False
        else:
            print(f"  {label:8s}  skipped (apply failed)")
            all_ok = False

    failed = [label for _, label, ok in results if not ok]
    print()
    if failed:
        print(f"FAILED: {', '.join(failed)}")
        sys.exit(1)
    elif all_ok:
        print("All routers configured successfully.")
        print()
        print(f"gRPC is now listening on port {GRPC_PORT} (no TLS) on all routers.")
        print("Next: start Telegraf with gNMI input plugin to begin collecting telemetry.")
    else:
        print("Some routers may have incomplete configuration. Check output above.")
        sys.exit(1)


if __name__ == '__main__':
    main()