#! /usr/bin/env python3
# UOSP firmware uploader
# Tested with Python 3.10.5

import opcode
from pickle import FALSE
from socket import timeout
from enum import Enum, IntEnum
from struct import pack
from time import sleep
from urllib import response
import serial
import sys
import re
import hashlib
import os

#from importlib.machinery import SourceFileLoader
#from crc16 import CRC16
#crc16 = SourceFileLoader("crc16", "/usr/bin/crc16").load_module();

def crc16(data: bytes):
    '''
    CRC-16 (CCITT) implemented with a precomputed lookup table
    '''
    table = [ 
        0x0000, 0x1021, 0x2042, 0x3063, 0x4084, 0x50A5, 0x60C6, 0x70E7, 0x8108, 0x9129, 0xA14A, 0xB16B, 0xC18C, 0xD1AD, 0xE1CE, 0xF1EF,
        0x1231, 0x0210, 0x3273, 0x2252, 0x52B5, 0x4294, 0x72F7, 0x62D6, 0x9339, 0x8318, 0xB37B, 0xA35A, 0xD3BD, 0xC39C, 0xF3FF, 0xE3DE,
        0x2462, 0x3443, 0x0420, 0x1401, 0x64E6, 0x74C7, 0x44A4, 0x5485, 0xA56A, 0xB54B, 0x8528, 0x9509, 0xE5EE, 0xF5CF, 0xC5AC, 0xD58D,
        0x3653, 0x2672, 0x1611, 0x0630, 0x76D7, 0x66F6, 0x5695, 0x46B4, 0xB75B, 0xA77A, 0x9719, 0x8738, 0xF7DF, 0xE7FE, 0xD79D, 0xC7BC,
        0x48C4, 0x58E5, 0x6886, 0x78A7, 0x0840, 0x1861, 0x2802, 0x3823, 0xC9CC, 0xD9ED, 0xE98E, 0xF9AF, 0x8948, 0x9969, 0xA90A, 0xB92B,
        0x5AF5, 0x4AD4, 0x7AB7, 0x6A96, 0x1A71, 0x0A50, 0x3A33, 0x2A12, 0xDBFD, 0xCBDC, 0xFBBF, 0xEB9E, 0x9B79, 0x8B58, 0xBB3B, 0xAB1A,
        0x6CA6, 0x7C87, 0x4CE4, 0x5CC5, 0x2C22, 0x3C03, 0x0C60, 0x1C41, 0xEDAE, 0xFD8F, 0xCDEC, 0xDDCD, 0xAD2A, 0xBD0B, 0x8D68, 0x9D49,
        0x7E97, 0x6EB6, 0x5ED5, 0x4EF4, 0x3E13, 0x2E32, 0x1E51, 0x0E70, 0xFF9F, 0xEFBE, 0xDFDD, 0xCFFC, 0xBF1B, 0xAF3A, 0x9F59, 0x8F78,
        0x9188, 0x81A9, 0xB1CA, 0xA1EB, 0xD10C, 0xC12D, 0xF14E, 0xE16F, 0x1080, 0x00A1, 0x30C2, 0x20E3, 0x5004, 0x4025, 0x7046, 0x6067,
        0x83B9, 0x9398, 0xA3FB, 0xB3DA, 0xC33D, 0xD31C, 0xE37F, 0xF35E, 0x02B1, 0x1290, 0x22F3, 0x32D2, 0x4235, 0x5214, 0x6277, 0x7256,
        0xB5EA, 0xA5CB, 0x95A8, 0x8589, 0xF56E, 0xE54F, 0xD52C, 0xC50D, 0x34E2, 0x24C3, 0x14A0, 0x0481, 0x7466, 0x6447, 0x5424, 0x4405,
        0xA7DB, 0xB7FA, 0x8799, 0x97B8, 0xE75F, 0xF77E, 0xC71D, 0xD73C, 0x26D3, 0x36F2, 0x0691, 0x16B0, 0x6657, 0x7676, 0x4615, 0x5634,
        0xD94C, 0xC96D, 0xF90E, 0xE92F, 0x99C8, 0x89E9, 0xB98A, 0xA9AB, 0x5844, 0x4865, 0x7806, 0x6827, 0x18C0, 0x08E1, 0x3882, 0x28A3,
        0xCB7D, 0xDB5C, 0xEB3F, 0xFB1E, 0x8BF9, 0x9BD8, 0xABBB, 0xBB9A, 0x4A75, 0x5A54, 0x6A37, 0x7A16, 0x0AF1, 0x1AD0, 0x2AB3, 0x3A92,
        0xFD2E, 0xED0F, 0xDD6C, 0xCD4D, 0xBDAA, 0xAD8B, 0x9DE8, 0x8DC9, 0x7C26, 0x6C07, 0x5C64, 0x4C45, 0x3CA2, 0x2C83, 0x1CE0, 0x0CC1,
        0xEF1F, 0xFF3E, 0xCF5D, 0xDF7C, 0xAF9B, 0xBFBA, 0x8FD9, 0x9FF8, 0x6E17, 0x7E36, 0x4E55, 0x5E74, 0x2E93, 0x3EB2, 0x0ED1, 0x1EF0
    ]
    
    crc = 0
    for byte in data:
        crc = (crc << 8) ^ table[(crc >> 8) ^ byte]
        crc &= 0xFFFF                                   # important, crc must stay 16bits all the way through
    return crc

# Global defines
SOH = 0x81
TYPE_DATA = 0x00
TYPE_ACK = 0x01
OPCODE_RESET = 0x00
OPCODE_ENABLE_BOOTLOADER = 0x2B

# Class that just read the file without verification
class IntelHexFile:
    def __init__(self, file) -> None:
        fileDescriptor = open(file, 'r')
        self.lines = fileDescriptor.readlines()
        self.count = 0
        fileDescriptor.close()

    def getNextLine(self):
        line = self.lines[self.count]
        self.count += 1
        return line

    def isEndOfFile(self):
        if len(self.lines) == self.count:
            return True
        else:
            return False

    def numberOfLines(self):
        return len(self.lines)

# Verifies if the microcontroller is in bootloader mode
def verifyIsBootloader(ser):
    # Send an enter to read back the command line
    ser.reset_input_buffer()
    ser.write(b'\r\r\n')

    # Around half a second for answer
    sleep(0.5)
    
    byte = ser.read(120)
    length = len(byte)

    if length > 10:
        # Search for the "> " at the end
        if hex(byte[-1]) == '0x20' and hex(byte[-2]) == '0x3e':
            return True

    return False

class GotDataState(IntEnum):
    INIT_STATE = 0,
    GOT_D = 1,
    GOT_FIRST_A = 2,
    GOT_T = 3,
    GOT_SECOND_A = 4,

# Enter the letter 'U' into the bootloader to enter the upload mode
def enterProgMode(ser):
    # Send the 'U' letter in the bootloader menu
    ser.write(b'U')
    sleep(0.2)
    ser.write(b'\r')
    sleep(0.2)
    ser.write(b'\n')

    sleep(3.5)
    answer = ser.read(120)

    if len(answer) < 10:
        print("Error, did not got an answer")
        return False

    # Search for "y " at the end of the answer
    if hex(answer[-1]) == '0x20' and hex(answer[-2]) == '0x79':
        # Write "JDSU" to unlock the upload mode
        ser.write(b'JDSU')

        currentState = GotDataState.INIT_STATE

        # Search for the word "DATA" (Waiting for DATA ...)
        i = 0
        while i < 400:
            i = i + 1
            byte = ser.read(1)

            if byte == b'D':
                if currentState == GotDataState.INIT_STATE:
                    currentState = GotDataState.GOT_D
                else:
                    print("Error, got an unexpected 'D'")
                    return FALSE
            
            if byte == b'A':
                if currentState == GotDataState.GOT_D:
                    currentState = GotDataState.GOT_FIRST_A
                if currentState == GotDataState.GOT_T:
                    currentState = GotDataState.GOT_SECOND_A
                    break

            if byte == b'T':
                if currentState == GotDataState.GOT_FIRST_A:
                    currentState = GotDataState.GOT_T
        
        if not currentState == GotDataState.GOT_SECOND_A:
            print("Could not find the string 'DATA', abort")
            return False
        else:
            return True

    return False

# Parse the address inside an Intel Hex line
def parseAddress(line):
    if len(line) < 7:  
        raise("Line is too short")
    
    addressString = line[3:7]
    return int(addressString, 16)

# Send the Intel Hex file in the UART
def sendFile(ser, intelHex):
    i = 0
    numberOfLines = intelHex.numberOfLines()
    previousAddress = 0
    addressJump = False
    while not intelHex.isEndOfFile():
        i += 1
        currentLine = intelHex.getNextLine()

        # Check the address. When the adress gap is too large,
        # the bootloader issues a write for the block before going
        # to the next block. And after the block write, the bootloader
        # sends an "OK" message back. Sending some bytes at the same
        # time as the Atmel will confuse the process and will
        # corrupt the data.
        address = parseAddress(currentLine)

        # Skip the first line adress check
        if i != 1:
            if address > previousAddress + 16:
                addressJump = True
            else:
                addressJump = False

        previousAddress = address

        if (i % 16 != 0) and not addressJump:
            ser.write(currentLine.encode())
        elif (i % 16 != 0) and addressJump:
            j = 0
            for byte in currentLine:
                if j == 8:
                    temp = ser.read(2)
                if j == 9:
                    temp = ser.read(2)
                    print("Got message : ", temp)
                ser.write(byte.encode())
                j = j + 1
        else:
            j = 0
            for byte in currentLine:
                # We receive a code at each block
                if j == len(currentLine) - 3 :
                    temp = ser.read(2)
                    print("Got message : ", temp)
                if j == len(currentLine) - 4 :
                    temp = ser.read(2)
                ser.write(byte.encode())
                j = j + 1

        if i % 16 == 0:
            completion = (i / numberOfLines) * 100
            print("Progress: ", completion, "%")

def sendEnableBootloader(serial, destination):
    buffer = [ SOH ]
    buffer.append(destination)
    buffer.append(0x00)
    buffer.append(TYPE_DATA)
    buffer.append(0x06) # Length 6 bytes
    buffer.append(0x00) # Length
    buffer.append(OPCODE_ENABLE_BOOTLOADER)
    buffer.append(0x04) # payload length = 4 bytes
    buffer.append(0x42) # B
    buffer.append(0x4F) # O
    buffer.append(0x4F) # O
    buffer.append(0x54) # T
    requestCRC = crc16(buffer) #crc16.CRC16(buffer)
    
    buffer.append(requestCRC & 0xFF)
    buffer.append(requestCRC >> 8)

    serial.write(buffer)
    return readAck(serial, destination)


def sendReset(serial, destination):
    buffer = [ SOH ]
    buffer.append(destination) # Destination
    buffer.append(0x00) # Source
    buffer.append(TYPE_DATA)
    buffer.append(0x02) # Length 2 bytes
    buffer.append(0x00) # Length
    buffer.append(OPCODE_RESET)
    buffer.append(0x00) # Length 0 bytes
    requestCRC = crc16(buffer) #crc16.CRC16(buffer)
    buffer.append(requestCRC & 0xFF)
    buffer.append(requestCRC >> 8)

    serial.write(buffer)
    return readAck(serial, destination)

def readAck(serial, destination):
    # Wait 50 ms for the answer
    sleep(50 / 1000)

    ack = serial.read(4)

    if len(ack) < 4:
        print("Wrong ack :", ack)
        return False
    else:
        if ack[0] == SOH and ack[1] == 0x00 and ack[2] == destination and ack[3] == TYPE_ACK:
            return True
        else:
            return False


def appToBootloader(serial, destination):
    if not sendEnableBootloader(serial, destination):
        print("Did not receive an ACK for the Enable bootloader request")
        return False

    # Wait 1 second before sending the reset
    sleep(1)
    if not sendReset(serial, destination):
        print("Did not receive the ACK for the reset request")
        return False

    sleep(3)
    # Wait and see if the device is in the bootloader mode
    if verifyIsBootloader(serial):
        print("In bootloader mode")
        return True
    else:
        print("The card does not seems to be in bootloader mode")
        return False


def bootloaderToApp(ser):
    # Send the 'R' letter in the bootloader menu
    ser.write(b'R')
    ser.write(b'R')
    sleep(90 / 1000)
    ser.write(b'\r')
    ser.write(b'\r')

    sleep(4.1)
    answer = ser.read(120)
    answerString = answer.decode('utf-8')

    #print("Bootloader to app: got answer:", answer)
    #print("Bootloader to app: got answer string :", answerString)

    # Search for "CLI disabled" in the answer
    regex = re.compile("\r\nCLI disabled\r\n")
    status = regex.match(answerString)
    if status:
        print("Now in application")
        return True
    else:
        print("Did not get the confirmation the UOS is now in application mode")
        return False

def verifyHexFile(filename):
    md5Filename = filename + ".md5"
 
    try:
        file = open(filename, "rb")
    except:
        print("Could not open file : ", filename)
        exit()

    sum = hashlib.md5(file.read()).hexdigest()

    try:
        checksumFile = open(md5Filename, 'r')
    except:
        print("Could not open file : ", md5Filename)
        exit()
    
    line = checksumFile.readline()
    checksumInFile = line.split(' ')

    if sum == checksumInFile[0]:
        return True
    else:
        return False

if __name__ == "__main__":
    print("File upload for UOSP bootloader")
    if len(sys.argv) < 3:
        print("Missing argument [Board address] [file to upload]")
        print("Format: uos_start_upgrade 1 /otu/release/current/tmp/Osmd_chksum.hex")
        exit(1)
    
    platform_name = os.getenv('CFG_MKT_PLATFORM_NAME')
    if platform_name:  
        if 'FTH-9000' in platform_name:
            com = "/dev/ttyRS485"
            print('FTH-9000 {}'.format(com))
        elif 'FTH-7000' in platform_name:
            com = "/dev/ttyRS485"
            print('FTH-7000 {}'.format(com))      
        else:
            com = "/dev/ttyS0"
            print('otu-5000 {}'.format(com))   
    else:
        platform_name = os.getenv('CFG_PLATFORM_NAME')
        if platform_name:
            if 'Otu5000' in platform_name:
                com = "/dev/ttyS0"
                print('otu-5000 {}'.format(com))
            else:
                com = "/dev/ttyS4"
                print('otu-8000e {}'.format(com))
        else:
            com = "/dev/ttyS4"
            print('otu-8000e {}'.format(com))

    #com = "/dev/ttyS0"
    #com = "COM3"
    #com = sys.argv[1]
    #address = 0x01
    address = int(sys.argv[1])
    #file = "Osmd_chksum.hex"
    file = sys.argv[2]

    print("Using COM  :", com)
    print("Using addr :", address)
    print("Using file :", file)
    
    if not verifyHexFile(file):
        print("The checksum value does not match the file value")
        exit(1)

    intelHex = IntelHexFile(file)
    with serial.Serial(com, 2400, timeout = 1) as ser:
        if not verifyIsBootloader(ser):
            print("Not in bootloader, switch to bootloader")
            if appToBootloader(ser, address):
                print("Now in bootloader, continue")
            else:
                print("Could not switch to bootloader, abort")
                exit(1)
        else:
            print("Inside the bootloader, try to enter prog mode ...")

        if not enterProgMode(ser):
            print("Error, not in prog mode")
            print('Quitting')
            exit(1)
        else:
            print("Inside prog mode, now send file ...")

        sendFile(ser, intelHex)
        print("File upload complete")

        # Flush the receive buffer after the update
        sleep(4)
        ser.reset_input_buffer()
        if not verifyIsBootloader(ser):
            print("Not in bootloader after update")

        if bootloaderToApp(ser):
            print("DONE")
            exit(0)
        else:
            print("The application may not have started")
            exit(1)
