From 172224cfdb9b05800e9355a714a151a8da12d638 Mon Sep 17 00:00:00 2001
From: Eljakim Herrewijnen <eljakim.herrewijnen@gmail.com>
Date: Wed, 17 Jan 2024 23:11:23 +0100
Subject: [PATCH] Initial

---
 Readme.md   |  23 ++++
 rigol_kg.py | 388 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 411 insertions(+)
 create mode 100644 Readme.md
 create mode 100644 rigol_kg.py

diff --git a/Readme.md b/Readme.md
new file mode 100644
index 0000000..781d2a1
--- /dev/null
+++ b/Readme.md
@@ -0,0 +1,23 @@
+# Rigol 
+Script from the EEVBlog, that will activate all options on a rigol scope.
+
+## Usage:
+```
+python3 rigol_kg.py 192.168.1.142
+```
+
+I had to force enable overwriting the key:
+```python
+def main(ip_addr):
+    [..]
+    prev_key = False
+
+    [..]
+    # ELHER necessary due to different key on my scope.
+    if len(new_key) < 388:
+        new_key += (388 - len(new_key)) * b"\xaa"
+
+    [..]
+```
+
+And reboot the scope and run again. Maybe it is not necessary and the reboot fixed the key?
\ No newline at end of file
diff --git a/rigol_kg.py b/rigol_kg.py
new file mode 100644
index 0000000..cf7f1ee
--- /dev/null
+++ b/rigol_kg.py
@@ -0,0 +1,388 @@
+import base64
+import binascii
+import os.path
+import re
+import socket
+import struct
+import sys
+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 = ['BW1T2', 'BW1T3', 'BW1T5', 'BW2T3', 'BW2T5',
+           'BW3T5', 'MSO', '2RL', '5RL', 'BND', 'COMP',
+           'EMBD', 'AUTO', 'FLEX', 'AUDIO', 'SENSOR',
+           'AERO', 'ARINC', 'DG', 'JITTER', 'MASK',
+           'PWR', 'DVM', 'CTR', 'EDK', '4CH', 'BW07T1',
+           'BW07T2', 'BW07T3', 'BW07T5']
+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, 0x1770078, 0x87888798])
+
+
+def decrypt_xxtea1(buf):
+    dec = xxtea.decrypt(buf, KEY1, padding=False)
+    return dec
+
+
+def decrypt_xxtea2(buf):
+    dec = xxtea.decrypt(buf, KEY2, padding=False)
+    return dec
+
+
+def encrypt_xxtea1(buf):
+    delta = len(buf) % 4
+
+    if delta:
+        buf += b'\x00' * (4 - delta)
+
+    enc = xxtea.encrypt(buf, KEY1, padding=False)
+    return enc
+
+
+def encrypt_xxtea2(buf):
+    delta = len(buf) % 4
+
+    if delta:
+        buf += b'\x00' * (4 - delta)
+
+    enc = xxtea.encrypt(buf, KEY2, 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)
+    buf = bytes(bb)
+    dig = sha256(buf).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
+
+
+BLOCK_HDR_FMT = '<IiIiI'
+
+
+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 write_block(buf, off, block, block_data):
+    block_len = len(block_data)
+    assert len(block['data']) == block_len
+    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, new_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 = {}
+
+    while off < full_size:
+        block, off = read_block(buf, off)
+        # print('%04X: id=%04d, data_sz=%04d, data=%s, crc32=%08X' % (
+        #     block['offset'], block['id'], len(block['data']), binascii.hexlify(block['data']).decode(), block['crc32']))
+        items[block['id']] = block
+
+    # print('last_off: 0x%04X' % off)
+
+    data = bytearray(cfram)
+    pub_key = items[4512]
+    write_block(data, 0x100, pub_key, new_key)
+
+    return bytes(data)
+
+
+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):
+    print('Reading CFRAM...')
+    cfram = bytearray()
+    i = 0
+
+    with tqdm(total=0x2000) as pb:
+        while i < 0x2000:
+            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 done.\n')
+
+    return bytes(cfram)
+
+
+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
+
+    model = m.group(1).decode()
+    ser = m.group(2).decode()
+    ver = m.group(3).decode()
+    mac = m.group(4).decode()
+
+    print('Model: %s\nSerial: %s\nVersion: %s\nMAC: %s\n' % (model, ser, ver, mac))
+
+    return model, ser
+
+
+def apply_new_key(ip_addr, new_cfram, new_key):
+    print('Applying new CFRAM...')
+    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(new_cfram)) as pb:
+        for i, b in enumerate(new_cfram):
+            exec_rigol_cmd(ip_addr, '/rigol/tools/fram01 -w %x %02x' % (i, b), need_res=False)
+            pb.update(1)
+
+    print('New CFRAM applied.\n')
+
+    exec_rigol_cmd(ip_addr, 'cp /rigol/data/Key.data /rigol/data/Key.data.bak', need_res=False)
+    print('Key.data backup created.')
+    exec_rigol_cmd(ip_addr, 'echo -n %s | base64 -d > /rigol/data/Key.data' % base64.b64encode(new_key).decode(), need_res=False)
+    print('New Key.data applied.')
+
+
+def check_fram_tool(ip_addr):
+    res = exec_rigol_cmd(ip_addr, 'md5sum /rigol/tools/fram')
+    res = res.split(' ')[0]
+
+    if res != FRAM_MD5:
+        print('Different /rigol/tools/fram hash. You have to recalc FRAM_OFFSETS!')
+        exit(-1)
+
+    print('/rigol/tools/fram is OK!\n')
+
+
+def deactivate_option(ip_addr, code, line):
+    print('Deactivating: %s...' % code, end=' ')
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.connect((ip_addr, 5555))
+    s.sendall(b':SYST:OPT:UNIN %s\n' % line)
+    print('deactivated.')
+    s.close()
+
+
+def activate_option(ip_addr, code, line):
+    print('Activating: %s...' % code, end=' ')
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.settimeout(1)
+    s.connect((ip_addr, 5555))
+    s.sendall(b':SYST:OPT:INST %s\n' % line)
+    s.sendall(b':SYST:OPT:STAT? %s\n' % code.encode())
+
+    try:
+        res = s.recv(2)
+
+        res = res.rstrip(b'\n')
+
+        if res == b'0':
+            print('not', end=' ')
+
+        print('activated.')
+    except TimeoutError:
+        print('unavailable.')
+    finally:
+        s.close()
+
+
+def get_unavail_options(ip_addr):
+    items = []
+    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 = [['Code', 'Status', 'Description']]
+    actives = []
+
+    for item in items:
+        row = item.split('$')
+        table.append(row)
+
+        if row[1] == 'Forever':
+            actives.append(row[0])
+
+    print(tabulate(table, headers='firstrow', tablefmt='fancy_grid'))
+
+    return list(set(OPTIONS) - set(actives))
+
+
+def main(ip_addr):
+    unavails = get_unavail_options(ip_addr)
+
+    check_fram_tool(ip_addr)
+
+    model, ser = read_rigol_model_serial(ip_addr)
+
+    opts = []
+    key_hex = None
+    prev_key = True
+
+    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))
+
+    # Set this to False to force overwrite the key
+    # prev_key = False
+    cfram = None
+    if not prev_key:
+        cfram = read_cfram_data(ip_addr)
+
+
+    new_key = b'brainpoolP256r1;%s' % key_hex
+    # ELHER necessary due to different key on my scope.
+    # if len(new_key) < 388:
+    #     new_key += (388 - len(new_key)) * b"\xaa"
+    new_key = encrypt_xxtea1(new_key)
+        
+    # new_key = 
+    # new_key = encrypt_xxtea2(b'brainpoolP256r1;%s' % key_hex)
+
+    if not prev_key:
+        new_cfram = replace_cfram_key(cfram, new_key)
+        apply_new_key(ip_addr, new_cfram, new_key)
+
+    model = model[:-2] + '00'
+
+    for opt in opts:
+        code = opt[0].encode()
+        activate_option(ip_addr, code.decode(), b'%s-%s@%s' % (model.encode(), code, opt[1]))
+
+    get_unavail_options(ip_addr)
+
+
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        print('Usage: python rigol_kg.py 192.168.1.1')
+    main(sys.argv[1])