#!/usr/bin/python3
import os, re, subprocess

# SHM User
shmUser = 'admin'

# Enable CDP
enableCDP = 'no'

# Backup directory
backupDir = '/home/cumulus/backup'

#Config for switch /etc/issue and /etc/banner.txt
banner = (
    '******************************************************************************',
    '*',
    '* NetApp Reference Configuration File (RCF)',
    '* Switch       : Mellanox MSN2100',
    '* Filename     : MSN2100-RCF-1.10-Cluster-HA-Breakout-LLDP',
    '* Release Date : 03-04-2023',
    '* Version      : 1.10 LLDP',
    '*',
    '* Port Usage:',
    '* Port 1      : 4x10G Breakout mode for Cluster+HA Ports, swp1s0-3',
    '* Port 2      : 4x25G Breakout mode for Cluster+HA Ports, swp2s0-3',
    '* Ports  3- 14: 40/100G for Cluster+HA Ports, swp3-14',
    '* Ports 15- 16: 100G Cluster ISL Ports, swp15-16',
    '*',
    '* NOTE:',
    '*   RCF manually sets swp1s0-3 link speed to 10000 and',
    '*   auto-negotiation to off for Intel 10G',
    '*   RCF manually sets swp2s0-3 link speed to 25000 and',
    '*   auto-negotiation to off for Chelsio 25G',
    '*',
    '*',
    '* IMPORTANT: Perform the following steps to ensure proper RCF installation:',
    '* - Copy the RCF file to /tmp',
    '* - Ensure the file has execute permission',
    '* - From /tmp run the file as sudo python3 <filename>',
    '*',
    '******************************************************************************\n',
)
#Config for adding a banner message
sshd_conf = {
    '#Banner none': 'Banner /etc/issue',
}
#Config for allowing passwordless use of cl-support by SHM
sudoers_cumulus_command_pre  = "echo '"
sudoers_cumulus_command_post = " ALL = NOPASSWD: /usr/cumulus/bin/cl-support, /usr/sbin/csmgrctl' | sudo EDITOR='tee -a' visudo -f /etc/sudoers.d/cumulus"

#Config for creating interfaces
add_interface = {
   'NCLU' : (
            'net add interface swp1 breakout 4x10G',
            'net add interface swp2 breakout 4x25G',
            'net add interface swp3-16 breakout 1x',
            'net commit',
           ),
   'NVUE' : (
            'nv set interface swp1 link breakout 4x',
            'nv set interface swp2 link breakout 4x',
            'nv set interface swp3-16 link breakout 1x',
            'nv config apply',
            'nv config save',
           ),
}

#Config for setting up interfaces
interface_config = {
   'NCLU' : (
            'net add bond cluster_isl bond slaves swp15,swp16',
            'net add bond cluster_isl bridge vids 1',
            'net add bond cluster_isl stp portnetwork',
            'net add bridge bridge ports swp1s0-3,swp2s0-3,swp3-14',
            'net add bridge bridge vids 1,17-18',
            'net add bridge bridge vlan-aware',
            'net add interface swp1s0-3 alias 10G Intra-Cluster Node',
            'net add interface swp2s0-3 alias 25G Intra-Cluster Node',
            'net add interface swp1s0-3,swp2s0-3 link autoneg off',
            'net add interface swp1s0-3 link speed 10000',
            'net add interface swp2s0-3 link speed 25000',
            'net add interface swp1s0-3 stp bpduguard',
            'net add interface swp1s0-3 stp portadminedge',
            'net add interface swp1s0-3 mtu 9216',
            'net add interface swp2s0-3 stp bpduguard',
            'net add interface swp2s0-3 stp portadminedge',
            'net add interface swp2s0-3 mtu 9216',
            'net add interface swp3-14 alias 40/100G Intra-Cluster Node',
            'net add interface swp3-14 stp bpduguard',
            'net add interface swp3-14 stp portadminedge',
            'net add interface swp3-14 mtu 9216',
            'net add interface swp15 alias Intra-Cluster Switch ISL Port swp15',
            'net add interface swp16 alias Intra-Cluster Switch ISL Port swp16',
            'net commit',
           ),
   'NVUE' : (
            'nv set bridge domain br_default type vlan-aware',
            'nv set bridge domain br_default vlan 1,17,18',
            'nv set interface cluster_isl bond member swp15,swp16',
            'nv set interface cluster_isl bridge domain br_default vlan 1',
            'nv set interface cluster_isl bridge domain br_default stp network on',
            'nv set interface cluster_isl type bond',
            'nv set interface swp15,swp16 type swp',
            'nv set interface swp3-14,swp1s0-3,swp2s0-3 bridge domain br_default vlan 1,17,18',
            'nv set interface swp3-14,swp1s0-3,swp2s0-3 type swp',
            'nv set interface swp1s0-3 description "10G Intra-Cluster Node"',
            'nv set interface swp2s0-3 description "25G Intra-Cluster Node"',
            'nv set interface swp1s0-3,swp2s0-3 link auto-negotiate off',
            'nv set interface swp1s0-3 link speed 10G',
            'nv set interface swp2s0-3 link speed 25G',
            'nv set interface swp1s0-3,swp2s0-3 link mtu 9216',
            'nv set interface swp1s0-3,swp2s0-3 bridge domain br_default stp admin-edge on',
            'nv set interface swp1s0-3,swp2s0-3 bridge domain br_default stp bpdu-guard on',
            'nv set interface swp15 description "Intra-Cluster Switch ISL Port swp15"',
            'nv set interface swp15-16 type swp',
            'nv set interface swp16 description "Intra-Cluster Switch ISL Port swp16"',
            'nv set interface swp3-14 description "40/100G Intra-Cluster Node"',
            'nv set interface swp3-14 link mtu 9216',
            'nv set interface swp3-14 bridge domain br_default stp admin-edge on',
            'nv set interface swp3-14 bridge domain br_default stp bpdu-guard on',
            'nv config apply',
            'nv config save',
           ),
}

#Config for disabling CDP /etc/default/lldpd
lldpd = (
    '# Add "-x" to DAEMON_ARGS to start SNMP subagent',
    '# NetApp disables CDP.  To enable CDP, add -c to',
    '# DAEMON_ARGS and issue <sudo systemctl restart lldpd.service>',
    'DAEMON_ARGS="-M 4 -x"\n',
)

#Config for enbling CDP /etc/default/lldpd
lldpd_cdp = (
    '# Add "-x" to DAEMON_ARGS to start SNMP subagent',
    '# NetApp disables LLDP.  To enable LLDP, remove -ll -c to',
    '# DAEMON_ARGS and issue <sudo systemctl restart lldpd.service>',
    'DAEMON_ARGS="-ll -c -M 4 -x"\n',
)

#Config for enabling lldp /etc/lldpd.d/lldpd.conf
lldpd_conf = (
    'configure lldp portidsubtype ifname',
    'configure lldp tx-interval 5',
    'configure lldp tx-hold 24\n',
)

#Config for enabling RoCE
roce_config = {
    'NCLU' : (
        'net add roce lossless',
        'net commit',
    ),
    'NVUE' : (
        'nv set qos roce mode lossless',
        'nv set qos mapping default-global trust both',
        'nv set qos pfc default-global switch-priority 2,5',
        'nv set qos pfc default-global cable-length 100',
        'nv set qos congestion-control default-global traffic-class 0,2,5 min-threshold 40000',
        'nv set qos congestion-control default-global traffic-class 0,2,5 max-threshold 200000',
        'nv set qos congestion-control default-global traffic-class 0,2,5 red enable',
        'nv set qos egress-scheduler default-global traffic-class 5 mode dwrr',
        'nv set qos egress-scheduler default-global traffic-class 5 bw-percent 43',
        'nv set qos egress-scheduler default-global traffic-class 0,2 mode dwrr',
        'nv set qos egress-scheduler default-global traffic-class 0,2 bw-percent 28',
        'nv set qos egress-scheduler default-global traffic-class 1,3,4,6,7 mode strict',
        'nv set qos traffic-pool default-lossy memory-percent 50',
        'nv set qos egress-queue-mapping default-global switch-priority 2 traffic-class 2',
        'nv set qos egress-queue-mapping default-global switch-priority 3 traffic-class 0',
        'nv set qos egress-queue-mapping default-global switch-priority 5 traffic-class 5',
        'nv set qos egress-queue-mapping default-global switch-priority 6 traffic-class 0',
        'nv config apply',
        'nv config save',
    )
}

#Config for allocating bandwidth to each queue and to add ECN, PFC /etc/cumulus/datapath/qos/qos_features.conf
traffic_conf = {
    'default_ecn_conf.egress_queue_list'          : 'default_ecn_conf.egress_queue_list = [0,2]',
    'default_egress_sched.egr_queue_0.bw_percent' : 'default_egress_sched.egr_queue_0.bw_percent = 28',
    'default_egress_sched.egr_queue_1.bw_percent' : 'default_egress_sched.egr_queue_1.bw_percent = 0',
    'default_egress_sched.egr_queue_2.bw_percent' : 'default_egress_sched.egr_queue_2.bw_percent = 28',
    'default_egress_sched.egr_queue_3.bw_percent' : 'default_egress_sched.egr_queue_3.bw_percent = 0',
    'default_egress_sched.egr_queue_4.bw_percent' : 'default_egress_sched.egr_queue_4.bw_percent = 0',
    'default_egress_sched.egr_queue_5.bw_percent' : 'default_egress_sched.egr_queue_5.bw_percent = 43',
    'default_egress_sched.egr_queue_6.bw_percent' : 'default_egress_sched.egr_queue_6.bw_percent = 0',
    'default_egress_sched.egr_queue_7.bw_percent' : 'default_egress_sched.egr_queue_7.bw_percent = 0',
    'default_ecn_conf.egress_queue_list'          : 'default_ecn_conf.egress_queue_list = [0,2,5]',
    'pfc.ROCE_PFC.cos_list'                       : 'pfc.ROCE_PFC.cos_list = [2,5]',
}
#Config for Class of service /etc/mlx/datapath/qos/qos_infra.conf
datapath_conf = {
    'egress_buffer.egr_queue_2.uc.service_pool'   : 'egress_buffer.egr_queue_2.uc.service_pool = 1',
    'egress_buffer.egr_queue_3.uc.service_pool'   : 'egress_buffer.egr_queue_3.uc.service_pool = 0',
    'egress_buffer.cos_2.mc.service_pool'         : 'egress_buffer.cos_2.mc.service_pool = 1',
    'egress_buffer.cos_3.mc.service_pool'         : 'egress_buffer.cos_3.mc.service_pool = 0',
    '#egress_buffer.egr_queue_2.uc.dynamic_quota' : 'egress_buffer.egr_queue_2.uc.dynamic_quota = ALPHA_8',
    'egress_buffer.egr_queue_3.uc.dynamic_quota'  : '#egress_buffer.egr_queue_3.uc.dynamic_quota = ALPHA_8',
    'cos_egr_queue.cos_0.uc'                      : 'cos_egr_queue.cos_0.uc = 0',
    'cos_egr_queue.cos_0.cpu'                     : 'cos_egr_queue.cos_0.cpu = 0',
    'cos_egr_queue.cos_1.uc'                      : 'cos_egr_queue.cos_1.uc = 0',
    'cos_egr_queue.cos_1.cpu'                     : 'cos_egr_queue.cos_1.cpu = 0',
    'cos_egr_queue.cos_2.uc'                      : 'cos_egr_queue.cos_2.uc = 2',
    'cos_egr_queue.cos_2.cpu'                     : 'cos_egr_queue.cos_2.cpu = 2',
    'cos_egr_queue.cos_3.uc'                      : 'cos_egr_queue.cos_3.uc = 0',
    'cos_egr_queue.cos_3.cpu'                     : 'cos_egr_queue.cos_3.cpu = 0',
    'cos_egr_queue.cos_4.uc'                      : 'cos_egr_queue.cos_4.uc = 0',
    'cos_egr_queue.cos_4.cpu'                     : 'cos_egr_queue.cos_4.cpu = 0',
    'cos_egr_queue.cos_5.uc'                      : 'cos_egr_queue.cos_5.uc = 5',
    'cos_egr_queue.cos_5.cpu'                     : 'cos_egr_queue.cos_5.cpu = 5',
    'cos_egr_queue.cos_6.uc'                      : 'cos_egr_queue.cos_6.uc = 0',
    'cos_egr_queue.cos_6.cpu'                     : 'cos_egr_queue.cos_6.cpu = 0',
    'cos_egr_queue.cos_7.uc'                      : 'cos_egr_queue.cos_7.uc = 0',
    'cos_egr_queue.cos_7.cpu'                     : 'cos_egr_queue.cos_7.cpu = 0',
}

#Config for enabling SNMP
nv_snmp_config_setup = {
    'NCLU' : (
        'net add snmp-server listening-address all vrf mgmt',
        'net add snmp-server mibs cumulus-status-mib',
        'net add snmp-server mibs cumulus-sensor-mib',
        'net commit',
    ),
    'NVUE' : (
        'nv set service snmp-server enable on',
        'nv set service snmp-server listening-address all vrf mgmt',
        'nv set service snmp-server mibs cumulus-status-mib',
        'nv set service snmp-server mibs cumulus-sensor-mib',
        'nv config apply',
        'nv config save',
    )
}

#enable the MIBS
nv_snmp_conf_add = {
    'NCLU' : (
        'net add snmp-server readonly-community "cshm1!" access any',
        'net commit',
    ),
    'NVUE' : (
        'nv set service snmp-server readonly-community "cshm1!" access any',
        'nv config apply --assume-yes',
        'nv config save',
    )
}

#Cleanup for restart
cleanup_config = {
   'NCLU' : (
            'net commit',
           ),
   'NVUE' : (
            # 'nv set system config apply ignore /etc/lldpd.d/lldpd.conf',
            # 'nv set system config apply ignore /etc/default/lldpd',
            # 'nv config apply',
            'nv config save',
           ),
}

#Switch initialization for NV to ignore /etc/lldpd.d/lldpd.conf and /etc/default/lldpd
initialization_config = (
    'nv set system config apply ignore /etc/lldpd.d/lldpd.conf',
    'nv set system config apply ignore /etc/lldpd.d/lldp-nvue.conf',
    'nv set system config apply ignore /etc/default/lldpd',
    'nv config apply',
    'nv config save',
)


#Default Cumulus version
cumulusVersion = 'NCLU'

def getCumulusVersion():
    result = ''
    lsb_release_path = "/etc/lsb-release"
    try:
        with open(lsb_release_path, mode='r', encoding='utf-8') as lsb_release_file:
            for line in lsb_release_file:
                match = re.search(r'DISTRIB_RELEASE=([\d.]+)', line)
                if match:
                    # found the release. Remove minor version number
                    version = re.sub(r'.\d+$', '', match.group(1))
                    if version:
                       result = version

    except FileNotFoundError:
        print("No", lsb_release_path, "file found")

    return result

if __name__== "__main__":
    # Get the version
    lsb_release = getCumulusVersion()
    if lsb_release:
        if lsb_release == '4.4':
           cumulusVersion = 'NCLU'
        elif lsb_release == '5.0':
           cumulusVersion = 'NVUE'
        elif lsb_release == '5.1':
           cumulusVersion = 'NVUE'
        else:
           cumulusVersion = 'NVUE'

    print('***Configure using', cumulusVersion, 'CLI***')
    
    #Initialization step
    if cumulusVersion == 'NVUE':
        print('***Step 0: Switch initialization***')
        for initilization_setp in initialization_config:
            os.system(initilization_setp)

    #Add issue and ssh banner
    print('***Step 1: Updating the MOTD file***')
    issue_file_path = '/etc/issue'
    os.chmod(f'{issue_file_path}', 0o777)
    #Write the file
    with open(issue_file_path, mode='wt', encoding='utf-8') as issue_file:
        issue_file.seek(0)
        issue_file.truncate()
        issue_file.write('\n'.join(banner))
    os.chmod(f'{issue_file_path}', 0o644)

    #Register the issue file with the banner directive in /etc/ssh/sshd_config
    print('***Step 2: Registering banner message***')
    sshd_config_dir_path = '/etc/ssh/'
    sshd_config_file_path = sshd_config_dir_path + 'sshd_config'
    #Save a copy of the file before changing it
    os.system(f'sudo mkdir -p {backupDir}{sshd_config_dir_path}')
    os.system(f'sudo cp {sshd_config_file_path} {backupDir}{sshd_config_file_path}_original')
    #Change the permission of the file to write
    os.chmod(f'{sshd_config_file_path}' , 0o777)
    with open(sshd_config_file_path, mode='r+', encoding='utf-8') as shd_config_file:
        data = shd_config_file.read()
        #Change the values
        for key, value in sshd_conf.items():
            data = re.sub(f'.*{key}.*', value, data)

        shd_config_file.seek(0)
        shd_config_file.truncate()
        shd_config_file.write(data)
    #Change the file permission back to original
    os.chmod(f'{sshd_config_file_path}', 0o644)

    #Add NOPASSWD directive to /etc/sudoers.d/cumulus for passwordless use of cl-support command by SHM
    print('***Step 3: Ensuring passwordless use of cl-support command by SHM***')
    sudoers_cumulus_command = sudoers_cumulus_command_pre + shmUser + sudoers_cumulus_command_post
    subprocess.call(sudoers_cumulus_command, stdout=subprocess.DEVNULL, shell=True)

    #Disable apt-get
    print('***Step 4: Disabling apt-get')
    os.chmod('/usr/bin/apt-get', 0o000)

    #Add interface non-breakout and breakout
    print('***Step 5: Creating the interfaces***')
    for config in add_interface[cumulusVersion]:
        os.system(config)

    # Add config and save
    print('***Step 6: Adding the interface config***')
    for int_conf in interface_config[cumulusVersion]:
        os.system(int_conf)

    #Configure lldp
    print('***Step 7: Configure lldpd***')
    lldp_dir_path = '/etc/default/'
    lldp_file_path = lldp_dir_path + 'lldpd'
    os.system(f'sudo mkdir -p {backupDir}{lldp_dir_path}')
    os.system(f'sudo cp {lldp_file_path} {backupDir}{lldp_file_path}_original')
    os.chmod(f'{lldp_file_path}', 0o777)
    if enableCDP == 'no':
        with open(lldp_file_path, mode='wt', encoding='utf-8') as lldp_file:
            lldp_file.seek(0)
            lldp_file.truncate()
            lldp_file.write('\n'.join(lldpd))
    else:
        with open(lldp_file_path, mode='wt', encoding='utf-8') as lldp_file:
            lldp_file.seek(0)
            lldp_file.truncate()
            lldp_file.write('\n'.join(lldpd_cdp))

    os.chmod(f'{lldp_file_path}', 0o644)

    print('***Step 8: Adding the lldp config***')
    lldpd_conf_file_path = '/etc/lldpd.d/lldpd.conf'
    os.system(f'sudo touch {lldpd_conf_file_path}')
    os.chmod(f'{lldpd_conf_file_path}', 0o777)
    with open(lldpd_conf_file_path, mode='wt', encoding='utf-8') as lldpd_conf_file:
        lldpd_conf_file.seek(0)
        lldpd_conf_file.truncate()
        lldpd_conf_file.write('\n'.join(lldpd_conf))
    os.system('sudo systemctl restart lldpd.service')

    # Add config and save
    print('***Step 9a: Adding the ROCE config***')
    for int_conf in roce_config[cumulusVersion]:
        os.system(int_conf)

    if cumulusVersion == 'NCLU':

        print('***Step 9b: Modifying RoCE Config***')
        traffic_conf_dir_path = '/etc/cumulus/datapath/qos/'
        traffic_conf_file_path = traffic_conf_dir_path + 'qos_features.conf'
        #Save a copy of the file before changing it
        os.system(f'sudo mkdir -p {backupDir}{traffic_conf_dir_path}')
        os.system(f'sudo cp {traffic_conf_file_path} {backupDir}{traffic_conf_file_path}_original')
        #Change the permission of the file to write
        os.chmod(f'{traffic_conf_file_path}' , 0o777)
        with open(traffic_conf_file_path, mode='r+', encoding='utf-8') as traffic_conf_file:
            data = traffic_conf_file.read()
            #Change the values
            for key, value in traffic_conf.items():
                data = re.sub(f'.*{key}.*', value, data)

            traffic_conf_file.seek(0)
            traffic_conf_file.truncate()
            traffic_conf_file.write(data)
        #Change the file permission back to original
        os.chmod(f'{traffic_conf_file_path}', 0o644)
        #Save a copy of the file before changing it
        datapath_conf_dir_path = '/etc/mlx/datapath/qos/'
        datapath_conf_file_path = datapath_conf_dir_path + 'qos_infra.conf'
        os.system(f'sudo mkdir -p {backupDir}{datapath_conf_dir_path}')
        os.system(f'sudo cp {datapath_conf_file_path} {backupDir}{datapath_conf_file_path}_original')
        #Change the permission of the file to write
        os.chmod(f'{datapath_conf_file_path}', 0o777)
        with open(datapath_conf_file_path, mode='r+', encoding='utf-8') as datapath_conf_file:
            data = datapath_conf_file.read()
            #Change the values
            for key, value in datapath_conf.items():
                data = re.sub(f'.*{key}.*', value, data)

            datapath_conf_file.seek(0)
            datapath_conf_file.truncate()
            datapath_conf_file.write(data)

        # #Change the file permission back to original
        os.chmod(f'{datapath_conf_file_path}', 0o644)


    # Add config and save
    print('***Step 10a: Enabling the SNMP config***')
    for snmp_conf in nv_snmp_config_setup[cumulusVersion]:
        os.system(snmp_conf)

    print('***Step 10b: Enabling the MIBS config***')
    for mibs_conf in nv_snmp_conf_add[cumulusVersion]:
        os.system(mibs_conf)

    #Cleanup for restart
    print('***Step 11: Cleanup***')
    for cleanup_config_line in cleanup_config[cumulusVersion]:
        os.system(cleanup_config_line)

    #Reboot the switch
    print('***Step 12: Reboot the switch***')
    os.system('sudo reboot')