CueHealth + Buspirate scripts

buspirate.py

import serial
import time
import struct

class BusPirateException(Exception):
    pass
    
class ModeSwitchError(BusPirateException):
    def __init__(self, mode, message="Mode switch failed"):
        self.mode = mode 
        self.message = message
        super().__init__(self.message)

        def __str__(self):
            return f'{self.message} (current mode: {self.mode})'

class BusPirate():
    def __init__ (self, port):
        self.ser = serial.Serial(port, 115200, timeout=0.1)
        self.write_command('\n')
        self.ser.reset_input_buffer()

    def write_command(self, command):
        command = command + '\n'
        self.ser.write(command.encode('utf-8'))
        # Wait for it to do its thing
        time.sleep(0.01)
        
    def query_command(self, command):
        self.ser.reset_input_buffer()
        self.write_command(command)

        line = '0'
        response = []

        while (line != b''):
            line = self.ser.readline()
            linedecoded = line.decode('utf-8')
            response.append(linedecoded.strip())
        return response

class BusPirateI2C(BusPirate):
# Parse the contents of the info message
    @staticmethod
    def parse_info(data):
        # Only return the things we actually need, we'll improve this functionality later
        info = {}
        info['cur_mode'] = data[-4]
        return info

# Enter I2C mode, enable power and pullups
    def enter_i2c (self):
        # Set mode as I2C, 100kHz
        self.write_command('m 4 3')

        # Validate mode switch
        info = self.parse_info(self.query_command('i'))

        if 'I2C' not in info['cur_mode']:
            raise ModeSwitchError(info['cur_mode'])

        # Enable power, pullup
        self.write_command('W')
        self.write_command('P')

    def dump_memory (self, dev_addr, start_addr, count):
        dump = []
        # Each address can access up to 2048 memory locations
        if start_addr + count > 2048:
            # Recursive function, start with partial starting block
            dump = self.dump_memory(dev_addr, start_addr, 2048 - start_addr)
            count = count - (2048 - start_addr)

            # Iterate through the remaining addresses until count runs out
            dev_addr_offset = 1
            while (count > 0):
                # If remaining count larger than one address, read 2048 and repeat for next block
                read_size = (2048 if count > 2048 else count)
                dump += self.dump_memory(dev_addr + 2 * dev_addr_offset, 0, read_size)
                dev_addr_offset += 1
                count -= 2048
        else:
            # Write command
            # note: ':02x' formats as hex, two digits, zero padded
            # message about addresses being dumped
            raw_data = self.query_command(f'[0x{(dev_addr):02X} 0x{(start_addr):02X} [0x{(dev_addr + 1):02X} r:{count}]]')

            # Get only the line that has the read data
            read_data = list(filter(lambda element: 'READ: ' in element, raw_data))

            # Split into individual memory contents
            dump = read_data[0].split("ACK ")

            # The first element will still have the READ command in it, so let's get rid of it
            dump[0] = dump[0].partition(' ')[2]

            # Finally, convert all to hex.
            dump = [bytes.fromhex(x.strip().replace('0x', '')) for x in dump]

        return dump

    def write_address (self, dev_addr, address, data):
        command = f'[0x{(dev_addr):02X} 0x{(address):02X} 0x{(data):02X}]'
        self.write_command(command)

    def write_multiple (self, dev_addr, address, data):
        for i in data:
            self.write_address(dev_addr, address, i)
            address += 1

dump-memory.py

# This script uses the encoding UTF-8
import time
import struct
from buspirate import BusPirateI2C, ModeSwitchError
from rich.console import Console
from rich.style import Style

#  ╭────────────────╮
#  │ User Constants │
#  ╰────────────────╯

# Buspirate serial device
BP_DEVICE = '/dev/tty.usbserial-AL05NDRG'

# Start address for chip
MEMORY_START = 0xA0
NUM_BYTES = 4096

# Rich styles
info_message = Style.parse("blue")
error_message = Style.parse("bold red")
success_message = Style.parse("green")

#  ╭───────────────╮
#  │ Main Function │
#  ╰───────────────╯

console = Console()
console.print(f'Dumping {NUM_BYTES} from device at address {MEMORY_START}')

# Get serial device
bp = BusPirateI2C(BP_DEVICE)

try:
    bp.enter_i2c()
except ModeSwitchError as e:
    console.print("Failed to switch modes!", style=error_message)
    console.print(e)

console.print("Beginning data dump", style=info_message)
dump = bp.dump_memory(MEMORY_START, 0, NUM_BYTES)

# TODO: if empty, print message

with open(f'hexdump-{int(time.time())}.hex', 'wb') as f:
    for item in dump:
        f.write(item)