#!/usr/bin/env python3
# Copyright (c) 2021 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

# pylint: disable=consider-using-f-string

from __future__ import absolute_import, division, print_function
import os, sys, time, argparse, datetime
from six.moves import range

####################################################################################
# misc helpers
####################################################################################

width = 0

CR   = '\r'
CEND = '\33[0m'
CRED = '\33[31m'
CGRN = '\33[32m'
CYEL = '\33[33m'
CDEL = '\33[1A'
CBLD = '\33[1m'
CUND = '\33[4m'

def printf(m):
   sys.stdout.write(m)
   sys.stdout.flush()
#end

msg = ''

def start( m ):
   global width

   width = 85

   global msg
   ts = str( datetime.datetime.now() )[ : 19 ]
   printf( '[%s] %s' % ( ts, m ) )
   msg = m

#end

def end( s=0, e=0, pad=0 ):
   if not isinstance( s, str ):
      s = ' OK ' if s == 0 else 'FAIL'
   dots = width - 35 - len( msg )
   printf( '.' * ( dots + pad ) + ' [ %s%s%s ]\n' %
     ( CRED if s == 'FAIL' else CGRN, s, CEND ) )
   if s == 'FAIL' and e:
      exit() # pylint: disable=consider-using-sys-exit
   #end
#end

def runcmd(c): 
   return os.popen(c).read()

def cli(c):
   c = 'printf "%s" | timeout 30 FastCli -Ap15' % c
   return runcmd(c)

def configure(c):
   c = '\n'.join(c)
   c = 'configure\n' + c + '\nend\n'
   return cli(c)

### get args ##################################################################

parser = argparse.ArgumentParser(description='BGP route scale script')

parser.add_argument( '-p', 
                     dest='peers',
                     action='store', 
                     help='number of peers',
                     default="10" )
parser.add_argument( '-r', 
                     dest='routes',
                     action='store', 
                     help='number of routes',
                     default="1M" )

parser.add_argument( '-d', 
                     dest='delete',
                     action='store_true',
                     help='delete configuration',
                     default=False )

args = parser.parse_args()

try:
   args.peers = int(args.peers)
   if args.peers > 100: 
      assert 0
except: #pylint: disable-msg=W0702
   print( "invalid number of peers - using 10 as default" )
#end

###############################################################################

INTF_IP_ADDR_FMT='%s.173.139.1'
INTF_IP_ADDR_BASE=201

cfg = cli('show running-config')

multiagent = 'service routing protocols model multi-agent' in cfg

cfg = cfg.split('!')

cfg = [ cg.strip() for cg in cfg ]

ifconfig = {}

for cg in cfg:
   if not cg.startswith('interface Eth'):
      continue
   intf = cg.split('\n')[0].split()[1]
   intf = intf.replace('Ethernet', 'et')
   ifconfig[intf] = cg
#end
  
if args.delete:
   start('unconfiguring interface IP addresses ...')
   for intf in ifconfig: # pylint: disable=consider-using-dict-items
      cg = ifconfig[intf]
      if INTF_IP_ADDR_FMT[3:] not in cg:
         continue
      configure([
         'interface %s' % intf,
         'swi',
         'no ip address'
      ])
   #end
   end()
   
   start( 'withdrawing BGP prefixes ...' )

   runcmd('gii set bgp route-scale delete peers all')

   end()

   start('unconfiguring BGP ...')

   configure(['no router bgp'])

   end()

   sys.exit()
#end

############################################################################### 

ifindex = 0

start( 'configuring interface IP addresses ...' )

configure(['ip routing'])

for intf in ifconfig: # pylint: disable=consider-using-dict-items

   entry = ifconfig[intf]
   if 'ip address ' in entry:
      continue
 
   ipaddr = ( INTF_IP_ADDR_FMT ) % ( INTF_IP_ADDR_BASE + ifindex )

   cfg = [
      'interface %s' % intf,
      'no switchport',
      'no ip address',
      'ip address %s 255.255.255.0' % ipaddr
   ]   

   configure(cfg)

   out = cli('bash ifconfig %s' % intf.replace('/', '_'))

   time.sleep(0.5)

   if ipaddr not in out: 
      configure([
         'interface %s' % intf,
         'swi',
         'no ip address'
      ])
      continue

   ifindex += 1

   if ifindex == args.peers:
      break

#end

if ifindex < args.peers:
   end(1)
   print( 'could not configure %s interface IP addresses' % args.peers )
   sys.exit()
#end

end()

###############################################################################

if multiagent:
   print( 'GateD/Rib single-agent mode is needed to inject BGP scale' )
   print( 'Remove the multi-agent configuration and rerun the script' )
   sys.exit()
#end

start( 'configuring BGP (%s peers) ...' % args.peers )

cmd = 'echo show running-config section bgp'
cfg = runcmd(cmd).split('\n')

bgp_asn = 0

for line in cfg:
   if line.startswith('router bgp '):
      bgp_asn = line.split('router bgp ')[1]
   #end
#end

if not bgp_asn:
   bgp_asn = '100'

peer_asn = bgp_asn + '1'

peers = []

for pindex in range( args.peers ):

   paddr = ( INTF_IP_ADDR_FMT + '1' ) % ( INTF_IP_ADDR_BASE + pindex )

   configure([
      'router bgp %s' % bgp_asn,
      'neighbor %s remote-as %s' % ( paddr, peer_asn ),
      'neighbor %s maximum-routes 0' % paddr
   ])

   peers += [paddr]

#end

peers = ' '.join(peers)

end()

###############################################################################

start( 'injecting %s BGP prefixes for each peer ...' % args.routes )

out = runcmd('gii set bgp route-scale %s peers %s' % (args.routes, peers))

end()

print()
print( '------------------------------------------------------------------------' )
print( ' GII output' )
print( '------------------------------------------------------------------------' )
print()
print( out )
time.sleep(5)
print()
print( '------------------------------------------------------------------------' )
print( ' ip route summary output' )
print( '------------------------------------------------------------------------' )
print()
print( cli( 'show ip route summary | grep bgp -A20' ) )

