import re
import json
import tabulate as tab
from texpect import *

KEY_DOWN = b'\x1b[B'

column_labels = {"atse2be":("be2Gci0FrameErrCount","be2Gci0FrameRxSorCount", "be2Gci0FramerStatus", "be2Gci0RetrainCount","be2Gci1FrameErrCount","be2Gci1FrameRxSorCount", "be2Gci1FramerStatus", "be2Gci1RetrainCount", "chip", "errorCount", "gci0BadAckFidCount", "gci0BadFidSorCount", "gci0RetrainBadCrcCount", "gci0RetrainCount", "gci1BadAckFidCount", 
                               "gci1BadFidSorCount", "gci1RetrainBadCrcCount", "gci1RetrainCount", "readCount", "unit", "writeCount"), 
                               "fpga_load":("chip", "class", "currentVersion", "pciBusList", "targetVersion", "unit", "updatePolicy", "updateStatus", "variant"),
                               "atse2fp":("chip", "name", "sourcePort", "unit", "value"),
                               "atse2vqf":("chip", "name", "sourcePort", "unit", "value"),
                               "vqf2qat":("chip", "name", "unit", "value"),
                               "vqf2bksim":("backplane", "name", "path", "port", "source", "value")}

##########################################################################
###                      begin script                                  ###
##########################################################################

blade = DUT("10.238.173.12", 7017, "jwhite", "default", "default", "KingstonDataTraveler 3.0PMAP", verbose=True)
blade.openConnection()

steps1 = (
	( b"root@", b"\r", b"reboot\r", 60 ),
	( b"Press <c>", b"", b"b", 200 ),
	( b"---/", b"", b"", 60 )
	)

for z in steps1: # Reboot DUT, get to the boot selection menu.
    try:
        match = texpect(blade.ssh_channel, z[0], z[1], z[2], timeout=z[3], verbose=blade.verbose)
        match.run()
    except TimeoutError:
        print(f"Timed out at steps1 step {z}")
        blade.closeConnection()
        exit()

boot_devs = re.findall(b'UEFI:.*?\\|',match.buffer)
boot_devs_desired = 0
for x, y in enumerate(boot_devs):
    if blade.boot_device in str(y): 
        boot_devs_desired = x
        break

if boot_devs_desired:
    print("Pressing down arrow to navigate to desired device...")
    for y in range(boot_devs_desired):
        blade.ssh_channel.send(KEY_DOWN)
        time.sleep(1)

steps2 = (
	( b"change the", "\r", b"\r", 60 ),
	( b"Writing obj", b"", b"", 120 ),
	( b"Press return", b"", b"\r", 300),
	( b"login:", b"", b"root\r", 300),
	( b"Password:", b"", b"default\r", 10),
	( b"UNIX password:", b"", b"default\r", 10),
	( b"New password:", b"", b"f5site02\r", 10),
	( b"Retype new password:", b"", b"f5site02\r", 10),
	( b"root@", b"", b"passwd root\r", 10),
	( b"New password:", b"", f"{blade.root_pass}\r", 10),
	( b"Retype new password:", b"", f"{blade.root_pass}\r", 10),
	( b"root@", b"", b'echo "QTP" > /etc/sysconfig/identity/blade\r', 10),
	( b"root@", b"", b"systemctl restart platform-services-deployment.service\r", 10),
	( b"c6xx 0000:", b"", b"dhclient\r", 500),
    ( b"root@", b"\r", f"""date -s '{time.strftime("%B %d %Y")}'""", 10),
    ( b"root@", b"\r", b"", 10)
	)

for w in steps2: # Perform OS install, configure in standalone mode, set password and date
    try:
        match = texpect(blade.ssh_channel, w[0], w[1], w[2], timeout=w[3], verbose=blade.verbose)
        match.run()
    except TimeoutError:
        print(f"Timed out at steps2 step {w}")
        blade.closeConnection()
        exit()

try_again=1
match = texpect(blade.ssh_channel, b"Setup complete", b"curl -fLkO https://sea.artifactory.f5net.com:443/artifactory/systems-eng/components-l4s/components-l4s-v3.4.65.tar.gz;tar -zxvf components-l4s-v3.4.65.tar.gz;cd components_l4s_3.4.65;./v6y_setup.py -f --use-local-components\r", b"", timeout=360, complications=[b"Traceback"], verbose=blade.verbose)
while try_again:
    try:
        if not match.run():
            try_again = 0
            print("Got containers!")
    except TimeoutError:
        pass
    except CaughtComplication:
        print("Failed to get containers. Waiting 20 seconds to try again.")
        time.sleep(20)
time.sleep(5) # A little nap for good measure

match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a localhost:9001 GET:lop/seeprom-platform\r", b"", verbose=blade.verbose)
match.run() # Get SEEPROM contents
for x in match.buffer[match.buffer.index(b"PcaPartNumber"):match.buffer.index(b"count :")].splitlines():
    y = re.findall(r'([a-zA-Z]+) +(.+)$', x.decode())
    if y:
        y = y.pop()
        blade.info[y[0]] = y[1]
blade.info = [x for x in blade.info.items()]

match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a localhost:9002 PUT:firmware/fpga chip=atse class=io unit=0\r", b"", verbose=blade.verbose)
match.run()
blade.setupresults.update({'atse0_load': [tuple([str(x) for x in list(json.loads(match.buffer[match.buffer.index(b"{"):match.buffer.index(b"}")+1].decode()).values())])]})
time.sleep(1)
match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a localhost:9002 PUT:firmware/fpga chip=vqf class=io unit=0\r", b"", verbose=blade.verbose)
match.run()
blade.setupresults.update({'vqf_load': [tuple([str(x) for x in list(json.loads(match.buffer[match.buffer.index(b"{"):match.buffer.index(b"}")+1].decode()).values())])]})
time.sleep(1)
match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a diag POST:measurements/atse/io/atse2fp duration=60\r", b"", timeout=180, verbose=blade.verbose)
match.run()
blade.testresults['atse2fp'] = re.findall(r"(atse|vqf|tam|nse) +([A-Z_]+) +(fp\d) +(\d+) +(\d+)", match.buffer[match.buffer.index(b"chip"):match.buffer.index(b"count :")].decode())
time.sleep(1)
match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a diag POST:measurements/fpga/io/atse2vqf\r", b"", verbose=blade.verbose)
match.run()
blade.testresults['atse2vqf'] = re.findall(r"(atse|vqf|tam|nse) +([A-Z_]+) +(ifh\d) +(\d+) +(\d+)", match.buffer[match.buffer.index(b"chip"):match.buffer.index(b"count :")].decode())
time.sleep(1)
match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a diag POST:measurements/vqf/io/vqf2qat\r", b"", verbose=blade.verbose)
match.run()
blade.testresults['vqf2qat'] = re.findall(r"(vqf) +([A-Z_]+) +(\d+) +(\d+)", match.buffer[match.buffer.index(b"chip"):match.buffer.index(b"count :")].decode())
time.sleep(1)
match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a diag POST:measurements/vqf/io/vqf2bksim\r", b"", verbose=blade.verbose, timeout=120)
match.run()
blade.testresults['vqf2bksim'] = re.findall(r"(bp\d+) +([A-Z_]+) +(\S+) +(\S+) +(\S+) +(\d+)", match.buffer[match.buffer.index(b"backplane"):match.buffer.index(b"count :")].decode())
time.sleep(1)
match = texpect(blade.ssh_channel, b"Exec Time", b"aclient -a diag measurements/atse/io/atse2be\r", b"", verbose=blade.verbose, timeout=120)
match.run()
blade.testresults.update({'atse2be': [tuple([str(x) for x in list(json.loads(match.buffer[match.buffer.index(b"{"):match.buffer.index(b"}")+1].decode()).values())])]})
# atse2be tests broken for now

blade_serial = ""
for x in blade.info:
    if "EndItemSerialNumber" in x:
        blade_serial = x[1]

with open(f"{blade_serial}_results.txt", 'w+') as save_file:
    print(f"Saving to file {blade_serial}_results.txt")
    save_file.write(f"{blade_serial} Test Results\n")
    save_file.write("Blade Information:\n")
    save_file.write(tab.tabulate(blade.info, tablefmt="outline"))
    save_file.write("\nATSE Load:\n")
    save_file.write(tab.tabulate(blade.setupresults['atse0_load'], headers=column_labels['fpga_load'], tablefmt="outline"))
    save_file.write("\nVQF Load:\n")
    save_file.write(tab.tabulate(blade.setupresults['vqf_load'], headers=column_labels['fpga_load'], tablefmt="outline"))
    save_file.write("\nATSE to Front Panel:\n")
    save_file.write(tab.tabulate(blade.testresults['atse2fp'], headers=column_labels['atse2fp'], tablefmt="outline"))
    save_file.write("\nATSE to VQF:\n")
    save_file.write(tab.tabulate(blade.testresults['atse2vqf'], headers=column_labels['atse2vqf'], tablefmt="outline"))
    save_file.write("\nVQF to QAT:\n")
    save_file.write(tab.tabulate(blade.testresults['vqf2qat'], headers=column_labels['vqf2qat'], tablefmt="outline"))
    save_file.write("\nVQF to Backplane Simulator:\n")
    save_file.write(tab.tabulate(blade.testresults['vqf2bksim'], headers=column_labels['vqf2bksim'], tablefmt="outline"))
    save_file.write("\nATSE to Bandwidth Engine:\n")
    save_file.write(tab.tabulate(blade.testresults['atse2be'], headers=column_labels['atse2be'], tablefmt="outline"))

steps3 = (
	( b"root@", b"\r", b"reboot\r", 60 ),
	( b"Press <c>", b"", b"b", 200 ),
	( b"---/", b"", b"", 60 )
	)

for z in steps3: # Reboot DUT, get to the boot selection menu.
    try:
        match = texpect(blade.ssh_channel, z[0], z[1], z[2], timeout=z[3], verbose=blade.verbose)
        match.run()
    except TimeoutError:
        print(f"Timed out at steps3 step {z}")
        blade.closeConnection()
        exit()

boot_devs = re.findall(b'UEFI:.*?\\|',match.buffer)
boot_devs_desired = 0
for x, y in enumerate(boot_devs):
    if blade.boot_device in str(y): 
        boot_devs_desired = x
        break

if boot_devs_desired:
    print("Pressing down arrow to navigate to desired device...")
    for y in range(boot_devs_desired):
        blade.ssh_channel.send(KEY_DOWN)
        time.sleep(1)

steps4 = (
	( b"change the", "\r", b"\r", 60 ),
	( b"Writing obj", b"", b"", 120 ),
	( b"Press return", b"", b"\r", 300),
	( b"login:", b"", b"root\r", 300),
	( b"Password:", b"", b"default\r", 10),
	( b"UNIX password:", b"", b"default\r", 10),
	( b"New password:", b"", b"f5site02\r", 10),
	( b"Retype new password:", b"", b"f5site02\r", 10),
	( b"root@", b"", b"passwd root\r", 10),
	( b"New password:", b"", f"{blade.root_pass}\r", 10),
	( b"Retype new password:", b"", f"{blade.root_pass}\r", 10)
	)

for w in steps4: # Perform OS install, set password and date
    try:
        match = texpect(blade.ssh_channel, w[0], w[1], w[2], timeout=w[3], verbose=blade.verbose)
        match.run()
    except TimeoutError:
        print(f"Timed out at steps4 step {w}")
        blade.closeConnection()
        exit()

blade.closeConnection()