#!/usr/bin/env python3
"""
Route Diversity Configuration Script
=====================================
Adds loopbacks, static routes, route-policies, and BGP redistribution
to R9K-01 through R9K-07 to create locally-originated routes that
produce meaningful RR Loc-RIB diffs between CORE-01 and CORE-02.

IS-IS topology (natural asymmetry — no metric tuning needed):
  CORE-01 → R9K-01, R9K-02, R9K-03, R9K-04, R9K-05
  CORE-02 → R9K-05, R9K-06, R9K-07
  R9K-04 ↔ R9K-06 (cross-link)
  R9K-05 dual-homed to both COREs

Overlapping prefixes from CORE-01-side and CORE-02-side routers produce
next-hop diffs because each RR picks the client with lowest IGP cost.

Address plan:
  Loopbacks:     10.110.{router_id}.{1,2}/32  (unique per router)
  Overlap LBs:   10.110.{100-103}.1/32        (shared across router pairs)
  Static routes: 10.111.{router_id}.0/24      (unique per router)
  Overlap statics: 10.111.{100-103}.0/24      (shared across router pairs)

Usage:
  python3 route_diversity_config.py              # apply all config
  python3 route_diversity_config.py --verify-only # just check current state
  python3 route_diversity_config.py --rollback    # remove all added config
"""

from ncclient import manager
import xml.etree.ElementTree as ET
import sys
import argparse

# YANG namespaces
IFMGR_NS   = 'http://cisco.com/ns/yang/Cisco-IOS-XR-ifmgr-cfg'
IPV4IO_NS  = 'http://cisco.com/ns/yang/Cisco-IOS-XR-ipv4-io-cfg'
STATIC_NS  = 'http://cisco.com/ns/yang/Cisco-IOS-XR-ip-static-cfg'
BGP_NS     = 'http://cisco.com/ns/yang/Cisco-IOS-XR-ipv4-bgp-cfg'
ISIS_NS    = 'http://cisco.com/ns/yang/Cisco-IOS-XR-clns-isis-cfg'
RPL_NS     = 'http://cisco.com/ns/yang/Cisco-IOS-XR-policy-repository-cfg'

# ──────────────────────────────────────────────────────────────────────
# Router definitions
# ──────────────────────────────────────────────────────────────────────

ROUTERS = {
    'R9K-01': {
        'mgmt': '10.100.0.1',
        'loopbacks': [
            ('Loopback10', '10.110.1.1', '255.255.255.255'),
            ('Loopback11', '10.110.1.2', '255.255.255.255'),
            ('Loopback100', '10.110.100.1', '255.255.255.255'),  # overlap with R9K-06
            ('Loopback103', '10.110.103.1', '255.255.255.255'),  # overlap with R9K-04, R9K-07
        ],
        'statics': [
            ('10.111.1.0', 24, 100),       # unique, tag=100 → LP=200
            ('10.111.100.0', 24, 100),      # overlap with R9K-06 (tag=200)
            ('10.111.103.0', 24, 100),      # overlap with R9K-04, R9K-07 (same tag)
        ],
    },
    'R9K-02': {
        'mgmt': '10.100.0.2',
        'loopbacks': [
            ('Loopback10', '10.110.2.1', '255.255.255.255'),
            ('Loopback11', '10.110.2.2', '255.255.255.255'),
            ('Loopback101', '10.110.101.1', '255.255.255.255'),  # overlap with R9K-07
        ],
        'statics': [
            ('10.111.2.0', 24, 100),
            ('10.111.101.0', 24, 100),      # overlap with R9K-07 (tag=300)
        ],
    },
    'R9K-03': {
        'mgmt': '10.100.0.3',
        'loopbacks': [
            ('Loopback10', '10.110.3.1', '255.255.255.255'),
            ('Loopback11', '10.110.3.2', '255.255.255.255'),
            ('Loopback102', '10.110.102.1', '255.255.255.255'),  # overlap with R9K-05
        ],
        'statics': [
            ('10.111.3.0', 24, 100),
            ('10.111.102.0', 24, 100),      # overlap with R9K-04 (tag=200)
        ],
    },
    'R9K-04': {
        'mgmt': '10.100.0.4',
        'loopbacks': [
            ('Loopback10', '10.110.4.1', '255.255.255.255'),
            ('Loopback11', '10.110.4.2', '255.255.255.255'),
            ('Loopback103', '10.110.103.1', '255.255.255.255'),  # overlap with R9K-01, R9K-07
        ],
        'statics': [
            ('10.111.4.0', 24, 200),        # tag=200 → LP=150
            ('10.111.102.0', 24, 200),       # overlap with R9K-03 (tag=100)
            ('10.111.103.0', 24, 100),       # overlap with R9K-01, R9K-07 (same tag)
        ],
    },
    'R9K-05': {
        'mgmt': '10.100.0.5',
        'loopbacks': [
            ('Loopback10', '10.110.5.1', '255.255.255.255'),
            ('Loopback11', '10.110.5.2', '255.255.255.255'),
            ('Loopback102', '10.110.102.1', '255.255.255.255'),  # overlap with R9K-03
        ],
        'statics': [
            ('10.111.5.0', 24, 100),
        ],
    },
    'R9K-06': {
        'mgmt': '10.100.0.6',
        'loopbacks': [
            ('Loopback10', '10.110.6.1', '255.255.255.255'),
            ('Loopback11', '10.110.6.2', '255.255.255.255'),
            ('Loopback100', '10.110.100.1', '255.255.255.255'),  # overlap with R9K-01
        ],
        'statics': [
            ('10.111.6.0', 24, 200),        # tag=200 → LP=150
            ('10.111.100.0', 24, 200),       # overlap with R9K-01 (tag=100)
        ],
    },
    'R9K-07': {
        'mgmt': '10.100.0.7',
        'loopbacks': [
            ('Loopback10', '10.110.7.1', '255.255.255.255'),
            ('Loopback11', '10.110.7.2', '255.255.255.255'),
            ('Loopback101', '10.110.101.1', '255.255.255.255'),  # overlap with R9K-02
            ('Loopback103', '10.110.103.1', '255.255.255.255'),  # overlap with R9K-01, R9K-04
        ],
        'statics': [
            ('10.111.7.0', 24, 300),         # tag=300 → LP=100
            ('10.111.101.0', 24, 300),        # overlap with R9K-02 (tag=100)
            ('10.111.103.0', 24, 100),        # overlap with R9K-01, R9K-04 (same tag)
        ],
    },
}

# ──────────────────────────────────────────────────────────────────────
# Route-policy (RPL text blob)
# ──────────────────────────────────────────────────────────────────────

ROUTE_POLICY_NAME = 'REDIST-TO-BGP'
ROUTE_POLICY_BODY = """\
route-policy REDIST-TO-BGP
  if tag is 100 then
    set local-preference 200
    set med 50
    set community (65020:100) additive
    pass
  elseif tag is 200 then
    set local-preference 150
    set med 100
    set community (65020:200) additive
    pass
  elseif tag is 300 then
    set local-preference 100
    set med 200
    set community (65020:300) additive
    pass
  else
    set local-preference 100
    pass
  endif
end-policy
"""

# ──────────────────────────────────────────────────────────────────────
# XML builders
# ──────────────────────────────────────────────────────────────────────

def loopback_xml(name, addr, mask):
    """Create a loopback interface with an IPv4 address."""
    return f"""
<config>
  <interface-configurations xmlns="{IFMGR_NS}">
    <interface-configuration>
      <active>act</active>
      <interface-name>{name}</interface-name>
      <interface-virtual/>
      <ipv4-network xmlns="{IPV4IO_NS}">
        <addresses>
          <primary>
            <address>{addr}</address>
            <netmask>{mask}</netmask>
          </primary>
        </addresses>
      </ipv4-network>
    </interface-configuration>
  </interface-configurations>
</config>
"""


def static_route_xml(prefix, prefix_len, tag):
    """Create a static route to Null0 with a tag."""
    return f"""
<config>
  <router-static xmlns="{STATIC_NS}">
    <default-vrf>
      <address-family>
        <vrfipv4>
          <vrf-unicast>
            <vrf-prefixes>
              <vrf-prefix>
                <prefix>{prefix}</prefix>
                <prefix-length>{prefix_len}</prefix-length>
                <vrf-route>
                  <vrf-next-hop-table>
                    <vrf-next-hop-interface-name>
                      <interface-name>Null0</interface-name>
                      <tag>{tag}</tag>
                    </vrf-next-hop-interface-name>
                  </vrf-next-hop-table>
                </vrf-route>
              </vrf-prefix>
            </vrf-prefixes>
          </vrf-unicast>
        </vrfipv4>
      </address-family>
    </default-vrf>
  </router-static>
</config>
"""


def route_policy_xml(name, body):
    """Create/replace a route-policy (RPL text blob)."""
    return f"""
<config>
  <routing-policy xmlns="{RPL_NS}">
    <route-policies>
      <route-policy>
        <route-policy-name>{name}</route-policy-name>
        <rpl-route-policy>{body}</rpl-route-policy>
      </route-policy>
    </route-policies>
  </routing-policy>
</config>
"""


def isis_passive_xml(intf_name):
    """Add a loopback to IS-IS instance 1 (passive by default for loopbacks)."""
    return f"""
<config>
  <isis xmlns="{ISIS_NS}">
    <instances>
      <instance>
        <instance-name>1</instance-name>
        <interfaces>
          <interface>
            <interface-name>{intf_name}</interface-name>
            <running/>
            <interface-afs>
              <interface-af>
                <af-name>ipv4</af-name>
                <saf-name>unicast</saf-name>
                <interface-af-data/>
              </interface-af>
            </interface-afs>
          </interface>
        </interfaces>
      </instance>
    </instances>
  </isis>
</config>
"""


def bgp_redistribute_xml():
    """Configure redistribute connected + static with REDIST-TO-BGP policy."""
    return f"""
<config>
  <bgp xmlns="{BGP_NS}">
    <instance>
      <instance-name>default</instance-name>
      <instance-as>
        <as>0</as>
        <four-byte-as>
          <as>65020</as>
          <bgp-running/>
          <default-vrf>
            <global>
              <global-afs>
                <global-af>
                  <af-name>ipv4-unicast</af-name>
                  <enable/>
                  <connected-routes>
                    <route-policy-name>{ROUTE_POLICY_NAME}</route-policy-name>
                  </connected-routes>
                  <static-routes>
                    <route-policy-name>{ROUTE_POLICY_NAME}</route-policy-name>
                  </static-routes>
                </global-af>
              </global-afs>
            </global>
          </default-vrf>
        </four-byte-as>
      </instance-as>
    </instance>
  </bgp>
</config>
"""


# ──────────────────────────────────────────────────────────────────────
# Rollback XML builders (delete operations)
# ──────────────────────────────────────────────────────────────────────

NC_NS = 'urn:ietf:params:xml:ns:netconf:base:1.0'


def delete_loopback_xml(name):
    return f"""
<config>
  <interface-configurations xmlns="{IFMGR_NS}">
    <interface-configuration xmlns:nc="{NC_NS}" nc:operation="delete">
      <active>act</active>
      <interface-name>{name}</interface-name>
    </interface-configuration>
  </interface-configurations>
</config>
"""


def delete_static_route_xml(prefix, prefix_len):
    return f"""
<config>
  <router-static xmlns="{STATIC_NS}">
    <default-vrf>
      <address-family>
        <vrfipv4>
          <vrf-unicast>
            <vrf-prefixes>
              <vrf-prefix xmlns:nc="{NC_NS}" nc:operation="delete">
                <prefix>{prefix}</prefix>
                <prefix-length>{prefix_len}</prefix-length>
              </vrf-prefix>
            </vrf-prefixes>
          </vrf-unicast>
        </vrfipv4>
      </address-family>
    </default-vrf>
  </router-static>
</config>
"""


def delete_bgp_redistribute_xml():
    return f"""
<config>
  <bgp xmlns="{BGP_NS}">
    <instance>
      <instance-name>default</instance-name>
      <instance-as>
        <as>0</as>
        <four-byte-as>
          <as>65020</as>
          <bgp-running/>
          <default-vrf>
            <global>
              <global-afs>
                <global-af>
                  <af-name>ipv4-unicast</af-name>
                  <enable/>
                  <connected-routes xmlns:nc="{NC_NS}" nc:operation="delete"/>
                  <static-routes xmlns:nc="{NC_NS}" nc:operation="delete"/>
                </global-af>
              </global-afs>
            </global>
          </default-vrf>
        </four-byte-as>
      </instance-as>
    </instance>
  </bgp>
</config>
"""


def delete_isis_interface_xml(intf_name):
    return f"""
<config>
  <isis xmlns="{ISIS_NS}">
    <instances>
      <instance>
        <instance-name>1</instance-name>
        <interfaces>
          <interface xmlns:nc="{NC_NS}" nc:operation="delete">
            <interface-name>{intf_name}</interface-name>
          </interface>
        </interfaces>
      </instance>
    </instances>
  </isis>
</config>
"""


def delete_route_policy_xml(name):
    return f"""
<config>
  <routing-policy xmlns="{RPL_NS}">
    <route-policies>
      <route-policy xmlns:nc="{NC_NS}" nc:operation="delete">
        <route-policy-name>{name}</route-policy-name>
      </route-policy>
    </route-policies>
  </routing-policy>
</config>
"""


# ──────────────────────────────────────────────────────────────────────
# Configuration functions
# ──────────────────────────────────────────────────────────────────────

def nc_connect(mgmt_ip):
    """Open NETCONF session."""
    return manager.connect(
        host=mgmt_ip,
        port=830,
        username='webui',
        password='cisco',
        hostkey_verify=False,
        device_params={'name': 'iosxr'},
        timeout=30,
    )


def configure_router(label, cfg):
    """Apply full route-diversity config to a single router."""
    mgmt_ip = cfg['mgmt']
    print(f"\n{'─'*60}")
    print(f"  Configuring {label} ({mgmt_ip})")
    print(f"{'─'*60}")

    lb_names = [lb[0] for lb in cfg['loopbacks']]
    static_prefixes = [f"{s[0]}/{s[1]}" for s in cfg['statics']]
    print(f"  Loopbacks: {', '.join(lb_names)}")
    print(f"  Statics:   {', '.join(static_prefixes)}")

    try:
        with nc_connect(mgmt_ip) as m:
            # Phase 1: Route-policy (must exist before BGP references it)
            print(f"  → Creating route-policy {ROUTE_POLICY_NAME}...")
            m.edit_config(target='candidate', config=route_policy_xml(ROUTE_POLICY_NAME, ROUTE_POLICY_BODY))

            # Phase 2: Loopback interfaces
            for name, addr, mask in cfg['loopbacks']:
                print(f"  → Creating {name} ({addr})...")
                m.edit_config(target='candidate', config=loopback_xml(name, addr, mask))

            # Phase 3: Static routes
            for prefix, plen, tag in cfg['statics']:
                print(f"  → Static {prefix}/{plen} → Null0 tag={tag}...")
                m.edit_config(target='candidate', config=static_route_xml(prefix, plen, tag))

            # Phase 4: IS-IS passive on new loopbacks
            for name, _, _ in cfg['loopbacks']:
                print(f"  → IS-IS passive: {name}...")
                m.edit_config(target='candidate', config=isis_passive_xml(name))

            # Phase 5: BGP redistribution
            print(f"  → BGP redistribute connected + static...")
            m.edit_config(target='candidate', config=bgp_redistribute_xml())

            # Phase 6: Commit
            print(f"  → Committing...")
            m.commit()
            print(f"  ✓ {label} done.")
            return True

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


def rollback_router(label, cfg):
    """Remove all route-diversity config from a single router."""
    mgmt_ip = cfg['mgmt']
    print(f"\n{'─'*60}")
    print(f"  Rolling back {label} ({mgmt_ip})")
    print(f"{'─'*60}")

    try:
        with nc_connect(mgmt_ip) as m:
            # Remove BGP redistribution first (references the policy)
            print(f"  → Removing BGP redistribute...")
            try:
                m.edit_config(target='candidate', config=delete_bgp_redistribute_xml())
            except Exception as e:
                print(f"    (skip — may not exist: {e})")

            # Remove IS-IS interfaces
            for name, _, _ in cfg['loopbacks']:
                print(f"  → Removing IS-IS interface {name}...")
                try:
                    m.edit_config(target='candidate', config=delete_isis_interface_xml(name))
                except Exception as e:
                    print(f"    (skip: {e})")

            # Remove static routes
            for prefix, plen, _ in cfg['statics']:
                print(f"  → Removing static {prefix}/{plen}...")
                try:
                    m.edit_config(target='candidate', config=delete_static_route_xml(prefix, plen))
                except Exception as e:
                    print(f"    (skip: {e})")

            # Remove loopbacks
            for name, _, _ in cfg['loopbacks']:
                print(f"  → Removing {name}...")
                try:
                    m.edit_config(target='candidate', config=delete_loopback_xml(name))
                except Exception as e:
                    print(f"    (skip: {e})")

            # Remove route-policy
            print(f"  → Removing route-policy {ROUTE_POLICY_NAME}...")
            try:
                m.edit_config(target='candidate', config=delete_route_policy_xml(ROUTE_POLICY_NAME))
            except Exception as e:
                print(f"    (skip: {e})")

            print(f"  → Committing rollback...")
            m.commit()
            print(f"  ✓ {label} rolled back.")
            return True

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


def verify_router(label, cfg):
    """Check if route-diversity config is present on a router."""
    mgmt_ip = cfg['mgmt']

    try:
        with nc_connect(mgmt_ip) as m:
            # Check loopbacks
            filt_intf = f"""<filter>
              <interface-configurations xmlns="{IFMGR_NS}"/>
            </filter>"""
            r_intf = str(m.get_config(source='running', filter=filt_intf))

            found_lbs = []
            for name, _, _ in cfg['loopbacks']:
                if name in r_intf:
                    found_lbs.append(name)

            # Check route-policy
            filt_rpl = f"""<filter>
              <routing-policy xmlns="{RPL_NS}"/>
            </filter>"""
            r_rpl = str(m.get_config(source='running', filter=filt_rpl))
            has_policy = ROUTE_POLICY_NAME in r_rpl

            # Check BGP redistribute
            filt_bgp = f"""<filter>
              <bgp xmlns="{BGP_NS}">
                <instance><instance-name>default</instance-name></instance>
              </bgp>
            </filter>"""
            r_bgp = str(m.get_config(source='running', filter=filt_bgp))
            has_redist_connected = 'connected-routes' in r_bgp and ROUTE_POLICY_NAME in r_bgp
            has_redist_static = 'static-routes' in r_bgp and ROUTE_POLICY_NAME in r_bgp

            total_lbs = len(cfg['loopbacks'])
            lb_str = f"{len(found_lbs)}/{total_lbs}"
            pol = '✓' if has_policy else '✗'
            rc = '✓' if has_redist_connected else '✗'
            rs = '✓' if has_redist_static else '✗'
            ok = len(found_lbs) == total_lbs and has_policy and has_redist_connected and has_redist_static
            status = 'OK' if ok else 'INCOMPLETE'

            print(f"  {label:8s}  LBs={lb_str:5s}  Policy={pol}  Redist-C={rc}  Redist-S={rs}  [{status}]")

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


# ──────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description='Route Diversity Configuration for RR Diff Analysis')
    parser.add_argument('--verify-only', action='store_true', help='Only verify current state')
    parser.add_argument('--rollback', action='store_true', help='Remove all added config')
    args = parser.parse_args()

    print("Route Diversity Configuration Script")
    print("=" * 60)
    print(f"Targets: {len(ROUTERS)} routers ({', '.join(ROUTERS.keys())})")
    print()

    if args.verify_only:
        print("Verify-only mode")
        print('-' * 60)
        print(f"  {'Router':8s}  {'LBs':5s}  {'Policy':6s}  {'Redist-C':8s}  {'Redist-S':8s}  Status")
        for label, cfg in ROUTERS.items():
            verify_router(label, cfg)
        return

    if args.rollback:
        print("ROLLBACK mode — removing all route-diversity config")
        print('-' * 60)
        results = []
        for label, cfg in ROUTERS.items():
            ok = rollback_router(label, cfg)
            results.append((label, ok))

        failed = [l for l, ok in results if not ok]
        print()
        if failed:
            print(f"FAILED rollback: {', '.join(failed)}")
            sys.exit(1)
        else:
            print("All routers rolled back successfully.")
        return

    # Apply mode
    results = []
    for label, cfg in ROUTERS.items():
        ok = configure_router(label, cfg)
        results.append((label, ok))

    # Post-apply verification
    print(f"\n{'='*60}")
    print("Post-apply verification")
    print('=' * 60)
    print(f"  {'Router':8s}  {'LBs':5s}  {'Policy':6s}  {'Redist-C':8s}  {'Redist-S':8s}  Status")
    for label, cfg in ROUTERS.items():
        verify_router(label, cfg)

    failed = [l for l, ok in results if not ok]
    print()
    if failed:
        print(f"FAILED: {', '.join(failed)}")
        sys.exit(1)
    else:
        total_lbs = sum(len(c['loopbacks']) for c in ROUTERS.values())
        total_statics = sum(len(c['statics']) for c in ROUTERS.values())
        print(f"All routers configured successfully.")
        print(f"  {total_lbs} loopbacks + {total_statics} static routes created")
        print()
        print("Wait ~60s for BGP convergence and BMP collection, then verify:")
        print()
        print("  # Check new prefixes in OpenBMP")
        print("  docker exec -i obmp-psql psql -U openbmp -d openbmp -c \\")
        print("    \"SELECT prefix::text, COUNT(*) FROM ip_rib")
        print("     WHERE (prefix::text LIKE '10.110.%' OR prefix::text LIKE '10.111.%')")
        print("     AND iswithdrawn = false GROUP BY prefix ORDER BY prefix;\"")


if __name__ == '__main__':
    main()