# Copyright (c) 2020 P&E Microcomputer Systems, Inc
# All rights reserved.
# Visit us at www.pemicro.com
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# o Redistributions of source code must retain the above copyright notice, this list
#   of conditions and the following disclaimer.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Version 1.00

'''
 Register ADKP for device secured in Challenge/Response mode
 This script requires a wrapped ADKP.
 This script requires two parameters:
 -hardwareid: the debug hardware IP address, name, serial number, or port name (e.g. -hardwareid=10.0.4.17)
 -wrapped_adkp: the 256 byte hexadecimal wrapped ADKP value (e.g. -wrapped_adkp=0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF)
 This script has one optional parameter:
 -user_pwd: if a smart card is being used and is password protected, the password must be provided (e.g. -user_pwd=pwd123)
'''

from pemicroprobewrapper import pemicroProbeWrapper
from pemicrounitacmp import pemicroUnitAcmp
from ctypes import *
import sys
import os


''' Constants '''

VOLKANO_DLL = 'volkano/volkano.dll'


''' Functions '''

# Load DLL Library
dll_path = os.path.dirname(os.path.abspath(__file__)) + os.path.sep + VOLKANO_DLL
libc = cdll.LoadLibrary(dll_path)

# Initialize DLL external functions
libc.vlk_init.argtypes = [c_char_p]
libc.vlk_init.restype = c_ulong
libc.vlk_register_key.argtypes = [POINTER(c_ubyte),c_ulonglong,POINTER(c_ubyte),c_ulonglong,c_ubyte]
libc.vlk_register_key.restype = c_ulong

# Create internal functions
def volkanoInit(pwd):
    cwd = os.getcwd()
    os.chdir('volkano')
    return libc.vlk_init(c_char_p(pwd.encode('utf-8')))
    os.chdir(cwd)
    
def registerADKP(uid, wrappedKey):
    retValue = 0
    
    uid = uid.ljust(16,'0')
    uidList = [0]*8
    for i in range(0,8):
        uidList[i] = int(uid[i*2:i*2+2], 16)
    uidBuffer = bytearray(uidList)
    uidSize = c_ulonglong(8)
    
    wpedKeyList = [0]*256
    for i in range(0,256):
        wpedKeyList[i] = int(wrappedKey[i*2:i*2+2], 16)
    wpedKeyBuffer = bytearray(wpedKeyList)
    wpedKeySize = c_ulonglong(256)
    
    keyType = c_ubyte(2) #ADKP
    
    return libc.vlk_register_key((c_ubyte*8).from_buffer(uidBuffer), uidSize, (c_ubyte*256).from_buffer(wpedKeyBuffer), wpedKeySize, keyType)


''' Get parameters (e.g. -hardwareid=10.0.4.17 -wrapped_adkp=0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF) '''

print('')
print('Registering Device.')
print('')

hardwareid=''
wrapped_adkp=''
user_pwd=''
for i in range(1, len(sys.argv)):
    if sys.argv[i].split('=')[0] == "-hardwareid":
        if len(sys.argv[i].split('=')) > 1: # if hardware ID is given
            hardwareid = sys.argv[i].split('=')[1]

    elif sys.argv[i].split('=')[0] == "-wrapped_adkp":
        if len(sys.argv[i].split('=')) > 1: # if wrapped ADKP is given
            wrapped_adkp = sys.argv[i].split('=')[1]
            if len(wrapped_adkp) != 512:
                sys.exit('Error - Wrapped ADKP must be 256 bytes (512 characters) long, in hexadecimal format (e.g. -wrapped_adkp=0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF).')

    elif sys.argv[i].split('=')[0] == "-user_pwd":
        if len(sys.argv[i].split('=')) > 1: # if user password ID is given
            user_pwd = sys.argv[i].split('=')[1]
            
if (hardwareid == '') or (wrapped_adkp == ''):
    if hardwareid == '':
        print('Error - Please specify the debug hardware IP address, name, serial number, or generic port name (such as USB1).')
        print('Examples of specifying the HW Identifier :')
        print('')
        print('             -hardwareid=USB1                ; Opens First PEmicro HW detected on USB bus')
        print('             -hardwareid=PEM592731           ; Opens PEmicro HW with Serial Number PEM592731')
        print('             -hardwareid=10.0.4.17           ; Opens PEmicro HW with IP 10.0.4.17')
        print('             -hardwareid=MyCyclone           ; Opens PEmicro HW with assigned Name MyCyclone')
        print('             -hardwareid=ETHERNET1           ; Opens First PEmicro HW detected on Ethernet')
        print('')
        print('Current List of autodetected PEmicro Hardware :')
        print('')
        pemicroUnitAcmp().listPortsDescription()
    if wrapped_adkp == '':
        print('Error - Please specify the 256-byte hexadecimal wrapped ADKP value needed to register the ADKP.')
        print('Example of specifying the wrapped ADKP :')
        print('')
        print('             -wrapped_adkp=0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF')
        print('')
    sys.exit(1)


''' Power up DAP main system domain and debug domain. '''

myprobe = pemicroProbeWrapper(0x07,hardwareid)

if not myprobe.connectedToDebugHardware:
    print ('Error ... Unable to connect to PEMicro HW',hardwareid)
    sys.exit(1)
else:
    print ('Connecting to debug hardware ... ok.')    

if not myprobe.connectedToTarget:
    print ('Error ... Unable to establish communications with target MCU')
    sys.exit(1)
else:
    print ('Requesting Debug Module Power On via SWD ... ok.')


''' Initialize Volkano DLL '''

retResult = volkanoInit(user_pwd)
if retResult != 0:
    print ('Initializing ... ERROR !!!!!!!!!!!!!!!!!!!!! (Error code: ', retResult,')')
    sys.exit(1)
else:
    print ('Initializing ... ok.')


''' Register key '''

# Test that we can talk to the SDA AP by reading its ID value (should read 0x001C0040)
targetRead = myprobe.read_reg(0xFC)
if not ((targetRead==0x001C0040) and (myprobe.pemicroProbeLibrary.lastSwdStatus()==0x04)):
    print ('Reading SDA AP ID Register ... ERROR !!!!!!!!!!!!!!!!!!!!! . (', hex(targetRead),')')
    sys.exit(1)
else:
    print ('Reading SDA AP ID Register ... ok.')

# Get UID value
uid0 = myprobe.read_reg(0x70)
uid1 = myprobe.read_reg(0x74)
uid_int = (uid1 << 32) | uid0
uid = hex(uid_int).replace('0x','')
print ('Reading SDA AP UID0 and UID1 ... ok.')

# Register ADKP
retResult = registerADKP(uid, wrapped_adkp)
if retResult != 0:
    print ('Registering device ... ERROR !!!!!!!!!!!!!!!!!!!!! (Error code: ', retResult,')')
    sys.exit(1)
else:
    print ('Registering device ... ok.')


myprobe.close()
del myprobe
print ('Success.') 
sys.exit(0)