import argparse
from argparse import RawTextHelpFormatter
import base64
import binascii
import os.path
import re
import socket
import struct
import time
import zlib
import requests
from tqdm import tqdm
from hashlib import sha256
from tabulate import tabulate

import xxtea
from struct import pack
from ecdsa import BRAINPOOLP256r1, SigningKey, VerifyingKey

FRAM_MD5 = 'aadc292fe4063a7ac392e3c3dde51e84'
FRAM_OFFSETS = [3120, 3204, 3206, 3207]
RESP_PAT = re.compile(rb'admin:\n(.+)\nOK\n', re.MULTILINE | re.DOTALL)
SERIAL_PAT = re.compile(rb'RIGOL TECHNOLOGIES\$(.+?)\$(.+?)\$(.+?)\n\$(.+?)\$(.+?)\$(.+?)\$(.+?)$')
OPTIONS = [
           #Bandwidth upgrades
           'BW1T2', 'BW1T3', 'BW1T5',
           'BW2T3', 'BW2T5',
           'BW3T5',
           'BW6T10', 'BW6T20', 'BW6T30', 
           'BW07T1', 'BW07T2', 'BW07T3', 'BW07T5',
           'BW10T20', 'BW10T30',
           'BW15T25', 'BW15T35',
           'BW20T30',
           'BW25T35',

           #Options
           'MSO', '4CH', '2RL', '5RL', 'BND', 'COMP', 'EMBD', 'AUTO', 'FLEX', 'AUDIO', 'SENSOR', 'AERO',
           'ARINC', 'DG', 'JITTER', 'MASK', 'PWR', 'DVM', 'CTR', 'AWG',

           #Educational license
           #'EDK'
           ]
PRIV_PATH = 'priv.pem'
KEY1 = b''.join(pack('<I', x) for x in [0x03920001, 0x08410841, 0x18C32104, 0x318639C7])
KEY2 = b''.join(pack('<I', x) for x in [0x478AA887, 0x99A85895, 0x01770078, 0x87888798])
BLOCK_HDR_FMT = '<IiIiI'

class tcolors:
    ENDC = '\033[0m'
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def print_term(text, color, nocolortext="", end="\n", ret=False):
    
    if USE_COLORS:
        string = f"{color}{text}{tcolors.ENDC}{nocolortext}"
    else:
        string = f"{text}{nocolortext}"
    if ret:
        return string
    print(string, end=end, flush=True)

def encrypt_xxtea(buf, key, pl_zero):
    buf += b'\x00' * (pl_zero-2)
    delta = len(buf) % 4
    if delta:
        buf += b'\x00' * (4 - delta)
    enc = xxtea.encrypt(buf, key, padding=False)

    return enc


def sign_option(opt):
    bb = bytearray()
    bb.extend(opt['model'].encode())
    bb.extend(opt['serial'].encode())
    bb.extend(opt['option'].encode())
    bb.extend(opt['version'].encode())
    bb.append(0x00)
    bb.append(0x00)
    dig = sha256(bytes(bb)).digest()

    prev_key = True
    if not os.path.exists(PRIV_PATH):
        sk = SigningKey.generate(curve=BRAINPOOLP256r1, hashfunc=sha256)
        with open(PRIV_PATH, 'wb') as w:
            w.write(sk.to_pem())
        prev_key = False
    else:
        with open(PRIV_PATH) as f:
            sk = SigningKey.from_pem(f.read())

    sign = sk.sign_digest_deterministic(dig)
    vk = sk.verifying_key
    vkk = b'04' + vk.to_string().hex().upper().encode()
    vk = VerifyingKey.from_string(binascii.unhexlify(vkk), curve=BRAINPOOLP256r1)

    assert vk.verify_digest(sign, dig)
    return binascii.hexlify(sign).upper(), vkk, prev_key


def calc_crc32(buf):
    return zlib.crc32(buf) & 0xFFFFFFFF


def get_dw(buf, off):
    dw = struct.unpack_from('<I', buf, off)[0]
    return dw, off + 4


def get_dws(buf, off):
    dw = struct.unpack_from('<i', buf, off)[0]
    return dw, off + 4


def get_data(buf, off, size):
    block = buf[off:off + size]
    return block, off + size


def read_block(buf, off):
    start = off
    id_, id_neg, data_size, data_size_neg, crc32 = struct.unpack_from(BLOCK_HDR_FMT, buf, off)
    off += struct.calcsize(BLOCK_HDR_FMT)

    assert ((id_ + id_neg) == 0) and ((data_size + data_size_neg) == 0)

    block_data, off = get_data(buf, off, data_size)
    crc32_real = calc_crc32(block_data)

    assert crc32_real == crc32

    return {
        'offset': start,
        'id': id_,
        'data': block_data,
        'crc32': crc32_real
    }, off


def neg(val):
    return -1 * val


def update_cfram_in_memory(buf, off, block, block_data):
    block_len = len(block_data)
    assert len(block['data']) == block_len, "generated public key length ("+str(len(block_data))+") does not match stored public key length ("+str(len(block['data']))+")"
    crc32 = calc_crc32(block_data)
    struct.pack_into(BLOCK_HDR_FMT, buf, off + block['offset'], block['id'], neg(block['id']), block_len, neg(block_len), crc32)
    block_data_off = off + block['offset'] + struct.calcsize(BLOCK_HDR_FMT)
    buf[block_data_off:block_data_off + block_len] = block_data


def replace_cfram_key(cfram, key_hex, using_key):
    buf = cfram[0x100:]

    off = 0
    full_size, off = get_dw(buf, off)
    full_size_neg, off = get_dws(buf, off)

    assert (full_size + full_size_neg) == 0

    items = {}
    pub_key = None
    if DEBUG:
        print_term("\nFRAM Parameters", tcolors.OKGREEN)
        fkey = str(["0x{:08X}".format(int.from_bytes(using_key[i:i+4], byteorder='little')) for i in range(0, len(using_key), 4)])
        print_term("Using xxtea key", tcolors.OKBLUE, f"={fkey}")

    while off < full_size:
        block, off = read_block(buf, off)
        if DEBUG:
            print_term("  offset", tcolors.OKCYAN, f"={block['offset']:04X},", end="")
            print_term(" id", tcolors.OKCYAN, f"={block['id']:04d},", end="")
            print_term(" data_sz", tcolors.OKCYAN, f"={len(block['data']):04d},", end="")
            print_term(" data", tcolors.OKCYAN, f"={binascii.hexlify(block['data']).decode():s},", end="")
            print_term(" crc32", tcolors.OKCYAN, f"={block['crc32']:08X}")
        try:
            # Not all devices have public key at 1D (29). Searching for a correct pubkey
            dect = xxtea.decrypt(block['data'], using_key, padding=False)
            if dect.decode().startswith('brainpoolP256r1;'):
                pub_key = block
                dec = dect
        except ValueError:
            # Decrypt failed, try the next block
            pass
        items[block['id']] = block
    if DEBUG:
        print_term("last read offset", tcolors.OKBLUE, f"=0x{off:04X}")
    assert pub_key is not None, "public key not found. Try a different key"
    if DEBUG:
        print_term("using public key at ", tcolors.OKBLUE, end="")
        print_term("offset", tcolors.OKCYAN, f"=0x{pub_key['offset']:04X}, ", end="")
        print_term("id", tcolors.OKCYAN, f"={pub_key['id']:04d}")

    pl_zero = len(re.search(r"\x00+$", dec.decode())[0])
    assert pl_zero % 2 == 0 , "0 padding in public key is not a multiple of 2"

    new_key = encrypt_xxtea(b'brainpoolP256r1;%s' % key_hex, using_key, pl_zero)

    data = bytearray(cfram)
    update_cfram_in_memory(data, 0x100, pub_key, new_key)

    return bytes(data), new_key


def exec_rigol_cmd(ip_addr, cmd, need_res=True):
    while True:
        res = requests.post('http://%s/cgi-bin/changepwd.cgi' % ip_addr, data={'pass0': '', 'pass1': '; %s # "' % cmd})

        if res.status_code == 500:
            continue

        body = res.content
        m = RESP_PAT.match(body)

        if m is None and need_res:
            continue

        if need_res:
            grp = m.group(1)
            return grp.decode()
        else:
            return None


def read_cfram_data(ip_addr, size=0x800):
    print('Reading CFRAM...\n**** DO NOT DISCONNECT POWER OR DATA CABLE ****', end="", flush=True)
    cfram = bytearray()
    i = 0

    # Just the first 2k of fram are important for our tasks
    with tqdm(total=size) as pb:
        while i < size:
            cmd = '/rigol/tools/fram -r %0x' % i
            res = exec_rigol_cmd(ip_addr, cmd)

            if res is None:
                continue

            bb = binascii.unhexlify(res.replace(',', ''))
            cfram.extend(bb)
            i += 0x10
            pb.update(0x10)
    print('Reading CFRAM... **** COMPLETED ****')
    return bytes(cfram)


def read_pubkey_data(ip_addr):
    print('Reading PublicKey...', flush=True)
    data = exec_rigol_cmd(ip_addr, 'cat /rigol/data/Key.data | base64')
    print('Reading PUBLIC KEY...**** COMPLETED ****')
    return base64.b64decode(data)


def read_rigol_model_serial(ip_addr):
    res = requests.post('http://%s/cgi-bin/welcome.cgi' % ip_addr)
    body = res.content
    m = SERIAL_PAT.match(body)

    if m is None:
        return None, None, None, None

    model = m.group(1).decode()
    ser = m.group(2).decode()
    ver = m.group(3).decode()
    mac = m.group(4).decode()

    return model, ser, ver, mac


def apply_new_key(ip_addr, new_key, diff_cfram):
    print('Applying new CFRAM...', end="", flush=True)
    exec_rigol_cmd(ip_addr, 'echo -n -e \'\\x03\' > /tmp/byte1', need_res=False)
    exec_rigol_cmd(ip_addr, 'echo -n -e \'\\x3d\' > /tmp/byte2', need_res=False)
    exec_rigol_cmd(ip_addr, 'echo -n -e \'\\x5b\' > /tmp/byte3', need_res=False)
    exec_rigol_cmd(ip_addr, 'echo -n -e \'\\xe5\' > /tmp/byte4', need_res=False)

    exec_rigol_cmd(ip_addr, 'cp /rigol/tools/fram /rigol/tools/fram01', need_res=False)
    exec_rigol_cmd(ip_addr, 'chmod +x /rigol/tools/fram01', need_res=False)
    exec_rigol_cmd(ip_addr, 'dd if=/tmp/byte1 of=/rigol/tools/fram01 obs=1 seek=%d conv=notrunc' % FRAM_OFFSETS[0], need_res=False)
    exec_rigol_cmd(ip_addr, 'dd if=/tmp/byte2 of=/rigol/tools/fram01 obs=1 seek=%d conv=notrunc' % FRAM_OFFSETS[1], need_res=False)
    exec_rigol_cmd(ip_addr, 'dd if=/tmp/byte3 of=/rigol/tools/fram01 obs=1 seek=%d conv=notrunc' % FRAM_OFFSETS[2], need_res=False)
    exec_rigol_cmd(ip_addr, 'dd if=/tmp/byte4 of=/rigol/tools/fram01 obs=1 seek=%d conv=notrunc' % FRAM_OFFSETS[3], need_res=False)

    with tqdm(total=len(diff_cfram)) as pb:
        for i in diff_cfram:
            exec_rigol_cmd(ip_addr, '/rigol/tools/fram01 -w %x %02x' % (i['offset'], i['new']), need_res=False)
            pb.update(1)

    print('.', end="", flush=True)
    exec_rigol_cmd(ip_addr, 'cp -f /rigol/data/Key.data /rigol/data/Key.data.bak', need_res=False)
    print('.', end="", flush=True)
    exec_rigol_cmd(ip_addr, 'echo -n %s | base64 -d > /rigol/data/Key.data' % base64.b64encode(new_key).decode(), need_res=False)
    print('New CFRAM applied - REGEN **** COMPLETED ****')


def check_fram_tool(ip_addr):
    if DEBUG:
        print("checking /rigol/tools/fram...", end='', flush=True);
    res = exec_rigol_cmd(ip_addr, 'md5sum /rigol/tools/fram')
    res = res.split(' ')[0]

    assert res == FRAM_MD5, '\nDifferent /rigol/tools/fram hash. You have to recalc FRAM_OFFSETS!'
    if DEBUG:
        print('DEBUG COMPLETED!')


def activate_ssh(ip_addr):
    print('Activating SSH...', end='', flush=True)
    exec_rigol_cmd(ip_addr, '/usr/sbin/sshd', need_res=False)
    print('****SSH ACTIVATED****\nNow..Open PuTTY.')

def wait_for_rigol(ip_addr):
    print("Waiting for device to be back online", end="", flush=True)
    while True:
        try:
            with socket.create_connection((ip_addr, 5555), timeout=1):
                break
        except OSError:
            time.sleep(0.01)
        print(".", end="", flush=True)
    print("DONE - PREPARE FOR REBOOT")


def options_uninstall(ip_addr):
    print('UNINSTALLING Activated Options......')
    print('==================================================')
    print('....**DO NOT DISCONNECT POWER OR DATA CABLE** ....')
    print('==================================================\n', end='', flush=True)
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(1)
    s.connect((ip_addr, 5555))
    s.sendall(b':SYSTem:OPTion:UNINSTall\n')
    s.close()
    print('UNINSTALLING... **** COMPLETED ****')


def activate_option(ip_addr, code, line):
    print('==================================================')
    print('....** ATTEMPTING ACTIVATION AND LICENSING ** ....')
    print('....**DO NOT DISCONNECT POWER OR DATA CABLE** ....')
    print('==================================================')
    print('ACTIVATING:\n%s\n[%s]...' % (code, line.decode()), end=' ')
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(1)
    s.connect((ip_addr, 5555))
    s.sendall(b':SYSTem:OPTion:INSTall %s\n' % line)
    s.sendall(b':SYSTem:OPTion:STAT? %s\n' % code.encode())
    print('\n')

    try:
        res = s.recv(2)

        res = res.rstrip(b'\n')

        if res == b'0':
            print('**** NOT ****', end=' ')

        print('ACTIVATED AND LICENSED.')
    except TimeoutError:
        print('Unavailable Option.')
    finally:
        s.close()


def get_unavail_options(ip_addr):
    items = []
    if DEBUG:
        print('==================================================')
        print('             **** DEBUGGING MODE ****             ')
        print('==================================================')
        print(f"Getting data from {ip_addr}...", flush=True)
        print('==================================================')
        print('  **** DO NOT DISCONNECT POWER OR DATA CABLE **** ')    
        print('==================================================')
    while True:
        res = requests.post('http://%s/cgi-bin/options.cgi' % ip_addr)

        if res.status_code == 500:
            continue

        body = res.content.decode()
        items = body.split('#')
        break

    table = [['OPTION CODE', 'STATUS', 'DESCRIPTION']]
    actives = []

    for item in items:
        row = item.split('$')
        table.append(row)

        if row[1] == 'Forever':
            actives.append(row[0])

    return list(set(OPTIONS) - set(actives)), table


def reboot(ip_addr):
    print('REBOOTING... ')
    try:
        exec_rigol_cmd(ip_addr, 'reboot', need_res=False)
    except:
        pass

    time.sleep(3)

    wait_for_rigol(ip_addr)
    print('REBOOT COMPLETED\n**** DEVICE ONLINE ****')


def model_to_license_str(model):
    m_bstr = ''
    m_num = ''
    if DEBUG:
        print(f"Input model: {model}", end="", flush=True)
    m = re.match(r'([A-Za-z]+)(\d+)([A-Za-z]*)$', model)
    if m is not None:
        m_bstr = m.group(1)
        m_num = m.group(2)
        model = m_bstr + m_num
        model = model[:len(m_bstr)+1] + ( '0' * (len(m_num)-1))
    if DEBUG:
        print(f" | Output model: {model}")
    return model

def generate_diff_cfram(old=b'', new=b''):
    old_ba = bytearray(old)
    new_ba = bytearray(new)
    diffs = []
    for i in range(0, len(old_ba)):
        if old_ba[i] != new_ba[i]:
            diffs.append({'offset': i, 'old': old_ba[i], 'new': new_ba[i]})
    return diffs

def write_to_file(ftype="_", data=None):
    time_exp = int(time.time() * 100 )
    filename = f"rigol_{ftype}_{time_exp}.data"
    f = open(filename, "wb")
    f.write(data)
    f.close()
    print(f"Saved {ftype} to Current Working Directory, {filename}")

def read_file(filename):
    with open(filename, mode="rb") as f:
        return bytes(f.read())
    return None

def get_config_name(id_):
    names = {
        '04': 'CH4',
        '03': 'CH3',
        '02': 'CH2',
        '01': 'CH1',
        '82': '---',
        '0B': 'Rigol Scope',
        '12': 'Math2',
        '11': 'Math1',
        '13': 'Math3',
        '14': 'Math4',
        '16': 'REF 1-10',
        '2F': 'Network Config',
    }
    if names.get(id_):
        return names[id_]
    return id_

def print_bin_setup(data):
    crc1 = struct.unpack('<I', data[:4])[0]
    l = struct.unpack('<I', data[4:8])[0]
    crc2 = zlib.crc32(data[8:8 + l])
    assert crc1 == crc2, f"setup data crc error {crc1:x} {crc2:x}"
    if DEBUG:
        print_term("\nFRAM Settings", tcolors.OKGREEN)
    l += 8
    i = 8
    odata = []
    while i < l:
        offset = i + 0x800
        idx = data[i]
        i += 1
        uncomp = data[i]
        i += 1
        x1 = data[i]
        i += 1
        x2 = data[i]
        i += 1
        dlen = struct.unpack('<H', data[i:i+2])[0]
        i += 2
        clen = struct.unpack('<H', data[i:i+2])[0]
        i += 2
        alen = struct.unpack('<H', data[i:i+2])[0]
        i += 2
        x3 = struct.unpack('<H', data[i:i+2])[0]
        i += 2
        bd = data[i:i+clen]
        bdx = bd
        zdec = True if uncomp == 0 else False
        if zdec:
            try:
                zr = zlib.decompressobj()
                bd = zr.decompress(bd[4:], dlen)
                if len(bd) != dlen:
                    i += alen
                    print(f"read {len(bd):x} {zr.unused_data}")
                    continue
            except zlib.error as err:
                print(f"zlib error {err}")
                return
        cfgname = get_config_name(f"{idx:02X}")
        if DEBUG:
            # offset
            print_term("  offset", tcolors.OKCYAN, f"={offset:04X},", end="")
            print_term(" id", tcolors.OKCYAN, f"={idx:02X} {x3:02X} {x2:02X}{x1:02X},", end="")
            print_term(" zlib_dec", tcolors.OKCYAN, f"={zdec:b},", end="")
            print_term(" uncomp_sz", tcolors.OKCYAN, f"={dlen:04d},", end="")
            print_term(" comp_sz", tcolors.OKCYAN, f"={clen:04d},", end="")
            print_term(" total_sz", tcolors.OKCYAN, f"={alen:04d},", end="")
            print_term(" name", tcolors.OKCYAN, f"={cfgname:s},", end="")
            print_term(" data", tcolors.OKCYAN, f"={binascii.hexlify(bd).decode('utf-8'):s}")
        i += alen

def main():
    parser = argparse.ArgumentParser(
        description=f'{tcolors.OKCYAN}RIGOL Oscilloscope Tool{tcolors.ENDC} for MSO5/6/7/8XXX v2.10b by {tcolors.OKCYAN}asp{tcolors.ENDC}',
        epilog="*********************************************************************\n"
        "                   SCROLL UP FOR HELP INFORMATION\n"
        "*********************************************************************\n"
        "****************************** CREDITS ******************************\n"
        "ORIGINAL VERSION BY DrMefistO\n"
        "---------------------------------------------------------------------\n"
        "LOTS OF THANKS TO :\n\n"
        "tv84                - Paving the path from the start\n"
        "                      Responsible for rigol_kg2.py\n\n"
        "SMAS                - Laying out the path to activation in an\n"
        "                      Easy to understand way.\n\n"
        "Seppletronics       - For working out the edited version of the file\n\n"
        "BTO                 - Assisting with Assertion Error\n"
        "                    - Minor script modifications resulting in\n"
        "                    - rigol_kg2_3_000.py\n"
        "                    - Rigol_MSO_LicensingUtility_2.09b\n"
        "                    - Cosmetic changes to scripts\n\n"
        "traxpalicaru (TrAx) - Donated his MSO8000 for testing\n\n"
        "asp                 - HUGE Effort in script progression\n"
        "                    - Opened the door to MSO7000 and MSO8000\n"
        "                    - Continued to work on ongoing Assertion Error\n"
        "                    - Optimized the script for faster activation\n"
        "                    - Responsible for...\n"
        "                    - rigol_mso_util_2.02a.py\n"
        "                    - rigol_mso_util_2.03a.py\n"
        "                    - rigol_mso_util_2.04a.py\n"
        "                    - rigol_mso_util_2.05b.py\n"
        "                    - rigol_mso_util_2.06b.py\n"
        "                    - rigol_mso_util_2.07b.py\n"
        "                    - rigol_mso_util_2.08b.py\n"
        "                    - rigol_mso_util_2.09b.py\n\n"
        "Kyr                 - Pointing out offest issue 0x001D\n\n"
        "ALL OTHER MEMBERS   - That supplied feedback when needed\n\n"
        "DAVE JONES(EEVBLOG) - For doing all that he does\n"
        "*********************************************************************\n"
        "                   SCROLL UP FOR HELP INFORMATION\n"
        "*********************************************************************\n",
        formatter_class=RawTextHelpFormatter)
    parser.add_argument("ip_addr", help="Rigol IP-address")
    parser.add_argument("-d", "--debug", help="Enable debugging", action="store_true")

    # action options
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("-i", "--info", help="Print options status, model and serial then exit", action="store_true")
    action.add_argument("-e", "--ssh", help="Start sshd (rigol/rigol or root/Rigol201)", action="store_true")
    action.add_argument("-r", "--regen", help="Regenerate private key", action="store_true")
    action.add_argument("-a", "--activate", help="Activate with private key", action="store_true")
    action.add_argument("-f", "--save-fram", help="Save first 2k of FRAM to 'rigol_fram_XXXXXXXXXXXX.data'. For DEBUG purposes only", action="store_true")
    action.add_argument("-p", "--save-pk", help="Save SSH Pub Key to 'rigol_pubkey_XXXXXXXXXXXX.data'. For DEBUG purposes only", action="store_true")
    action.add_argument("-u", "--uninstall", help="Uninstall all options.", action="store_true")
    action.add_argument("--reboot", help="Reboots the device and waits for it to be back online", action="store_true")
    action.add_argument("--cfram-file", help="Dummy run based contents read from cfram file export", metavar="rigol_fram_XXXXXXXXXXXX.data")

    # parameters
    params = parser.add_argument_group("Parameters", "Different parameters")
    params.add_argument("-k", "--keyver", help="Key version", default=1, type=int, choices=[1,2])
    params.add_argument("-m", "--model", help="Model Name", default=None, metavar="MSO5074")
    params.add_argument("-s", "--serial", help="Serial Number", default=None, metavar="MS5A123456789")
    params.add_argument("-o", "--options", help="Options", default=None, metavar="AUTO,AERO,AUDIO")
    params.add_argument("--sys-vendor-file", help="sysvendor file ***NOT WORKING***", default=None, metavar="rigol_vendor_XXXXXXXXXXXX.data")
    params.add_argument("--no-reboot", help="Do not reboot scope after regenerating keys", action="store_true")
    params.add_argument("--with-config", help="Will get 8k from FRAM. This will be slower", action="store_true")
    params.add_argument("--no-color", help="Color output", action="store_true")

    args = parser.parse_args()
    dummy_run = False

    global DEBUG
    global USE_COLORS
    if args.debug:
        DEBUG=True
    else:
        DEBUG=False
    if args.no_color:
        USE_COLORS = False
    else:
        USE_COLORS = True
    
    if args.cfram_file is not None and os.path.isfile(args.cfram_file):
        if ((args.model is None) or (args.serial is None) or (args.options is None)) and args.sys_vendor_file is None:
            print("Dummy [--model,--serial and --options | --sys-vendor-file] are mandatory when running with cfram file")
            exit(-1)
        model = args.model
        ser = args.serial
        ver = 'DUMMY'
        mac = '0F-F1-C1-A1-BA-BE'
        unavails = args.options.split(',')
        k_model = model_to_license_str(model)
        dummy_run = True
    else:
        #Read model and serial
        model, ser, ver, mac = read_rigol_model_serial(args.ip_addr)

    if args.model is not None and not dummy_run:
        k_model = args.model
    else:
        k_model = model_to_license_str(model)

    print(tabulate([['Model', 'Serial', 'Version', 'MAC', 'Lic Model'],[model,ser,ver,mac,k_model]], headers='firstrow', tablefmt='fancy_grid'))

    if args.options is None and not dummy_run:
        #Get options from scope
        unavails, table = get_unavail_options(args.ip_addr)
        print(tabulate(table, headers='firstrow', tablefmt='fancy_grid'))

    #Read fram size
    fsize = 0x800
    if args.with_config:
        fsize = 0x2000

    if not dummy_run:
        if args.info:
            return

        if args.ssh:
            activate_ssh(args.ip_addr)
            return

        if args.uninstall:
            options_uninstall(args.ip_addr)
            reboot(args.ip_addr)
            return

        if args.save_fram:
            write_to_file("fram", read_cfram_data(args.ip_addr, fsize))
            return

        if args.save_pk:
            write_to_file("Public Key", read_pubkey_data(args.ip_addr))
            return

        if args.reboot:
            reboot(args.ip_addr)
            return

        #Check fram tool
        check_fram_tool(args.ip_addr)

    opts = []
    key_hex = None
    prev_key = True
    cfram = None

    # Delete previous priv.pem if --regen passed as an argument
    if args.regen and os.path.isfile('priv.pem'):
        os.remove('priv.pem')

    for option in unavails:
        opt_sign, key_hex, prev_key = sign_option({
            'model': model,
            'serial': ser,
            'option': option,
            'version': '1.0'
        })
        opts.append((option, opt_sign))

    opts = sorted(opts, key=lambda e: e[0])

    #Defaults to KEY1
    using_key = KEY1
    if args.keyver == 2:
        using_key = KEY2

    # --regen passed as an argument so we will regenerate a new private key and update cfram
    if args.regen:
        cfram = read_cfram_data(args.ip_addr, fsize)
    elif dummy_run:
        cfram = read_file(args.cfram_file)

    if cfram is not None:
        i = 7
        print_term("System Setup", tcolors.OKGREEN)
        print_term("  Language: ", tcolors.OKCYAN, f"{cfram[i+1]:X}" )
        print_term("  LoadLast: ", tcolors.OKCYAN, f"{cfram[i+2]:X}" )
        print_term("  PowerStatus: ", tcolors.OKCYAN, f"{cfram[i+3]:X}" )
        print_term("  GPIB: ", tcolors.OKCYAN, f"{cfram[i+4]:X}" )
        print_term("  Cycle count: ", tcolors.OKCYAN, (int.from_bytes(cfram[i+5:i+13], byteorder='big') & 0x0000FFFFFFFFFFFF))
        print_term("  Live count: ", tcolors.OKCYAN, int.from_bytes(cfram[i+13:i+17], byteorder='big'))
        print_term("  Keep Imp: ", tcolors.OKCYAN, f"{cfram[i+18]:X}" )
        print_term("  IPmode: ", tcolors.OKCYAN, f"{cfram[i+19]:X}" )
        if len(cfram) >= 0x2000:
            print_bin_setup(cfram[0x800:])

    if args.regen or dummy_run:
        # replace key in memory
        new_cfram, new_key = replace_cfram_key(cfram, key_hex, using_key)
        # generate diff between old and new cfram data
        diff_cfram = generate_diff_cfram(cfram, new_cfram)
        if not dummy_run:
            # apply new key to the oscilloscope
            apply_new_key(args.ip_addr, new_cfram, diff_cfram)
            if args.no_reboot:
                return
            # reboot the scope
            reboot(args.ip_addr)

    if args.activate or args.regen:
        for opt in opts:
            code = opt[0].encode()
            activate_option(args.ip_addr, code.decode(), b'%s-%s@%s' % (k_model.encode(), code, opt[1]))

    if not dummy_run:
        _, table = get_unavail_options(args.ip_addr)
        print(tabulate(table, headers='firstrow', tablefmt='fancy_grid'))


if __name__ == '__main__':
    main()