#!/usr/bin/python3
import re
import string
import hashlib

# adapted from https://www.eevblog.com/forum/testgear/unlocking-siglent-sds1104x-e-step-by-step/msg1973306/
def getkeys(scopeid, serialno, memdumpfile):
    """
    Parse a memory dump from a Siglent 1000X-E oscilloscope and return a dict containing
    license keys. The 'activebw' key is the one that is currently active in the 'scope 
    (e.g. if the value of '100M' is the same as the value of 'activebw',
    the oscilloscope is software locked to 100 MHz bandwidth)
    """
    if len(scopeid) == 16 and set(scopeid) <= set(string.hexdigits):
        scopeid = scopeid.lower().encode('utf-8')
    else:
        raise ValueError('Scope ID must be 16 hexadecimal characters (remove dashes).')
   
    if len(serialno) == 14 and set(serialno) <= set(string.ascii_letters + string.digits):
        serialno = serialno.upper().encode('utf-8')
    else:
        raise ValueError('Serial number must be 14 alphanumeric characters.')

    f = open(memdumpfile, 'rb')
    data = f.read()
    f.close()

    regex_bw = re.compile(scopeid + b'.*?'+ scopeid + b'.*?([0-9A-Z]{16}).*?([0-9A-Z]{16}).*?([0-9A-Z]{16}).*?([0-9A-Z]{16}).*?([0-9A-Z]{16})', re.DOTALL)
   
    key_bw = list([n.decode('utf-8') for n in re.findall(regex_bw, data)[0]])
   
    keys = {}
    key_labels = ('100M', '200M', '50M', '70M', 'activebw')
    keys.update(zip(key_labels, key_bw))

    return(keys)

# adapted from https://www.eevblog.com/forum/testgear/unlocking-siglent-sds1104x-e-step-by-step/msg3407174/
def genopts(serial):
    options = {}
    for opt in ('AWG', 'WIFI', 'MSO'):
       h = hashlib.md5((
          '5zao9lyua01pp7hjzm3orcq90mds63z6zi5kv7vmv3ih981vlwn06txnjdtas3u2'
          'wa8msx61i12ueh14t7kqwsfskg032nhyuy1d9vv2wm925rd18kih9xhkyilobbgy' +
          'SDS1000X-E\n'.ljust(32, '\x00') +
          opt.ljust(5, '\x00') +
          2*((serial + '\n').ljust(32, '\x00')) +
          '\x00'*16).encode('ascii')
       ).digest()
       key = ''
       for b in h:
          if (b <= 0x2F or b > 0x39) and (b <= 0x60 or b > 0x7A):
             m = b % 0x24
             b = m + (0x57 if m > 9 else 0x30)
          if b == 0x30: b = 0x32
          if b == 0x31: b = 0x33
          if b == 0x6c: b = 0x6d
          if b == 0x6f: b = 0x70
          key += chr(b)
       options[opt] = key.upper()
    return options

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Recover Siglent SDS1004X-E bandwidth/option licenses")
    parser.add_argument("--sid", dest="scope_id", type=str, required=False, help="Scope ID")
    parser.add_argument("--serial", dest="serial", type=str, required=True, help="Scope serial number")
    parser.add_argument("--dump", dest="memdump", type=str, required=False, help="Path to the memory dump from the scope")
    args = parser.parse_args()
    keys = dict()
    if args.scope_id and args.memdump:
        keys.update(getkeys(args.scope_id, args.serial, args.memdump))
    keys.update(genopts(args.serial))
    print(keys)