Rob Barnes | d6b58d5 | 2023-06-14 11:19:53 -0600 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
| 2 | |
| 3 | # Script for injecting SPDs into APCB_v3a binaries. |
| 4 | |
| 5 | import re |
| 6 | import argparse |
| 7 | from collections import namedtuple |
| 8 | from struct import * |
| 9 | |
| 10 | APCB_CHECKSUM_OFFSET = 16 |
| 11 | SPD_ENTRY_MAGIC = bytes.fromhex('0200480000000000') |
| 12 | SPD_SIZE = 512 |
| 13 | EMPTY_SPD = b'\x00' * SPD_SIZE |
| 14 | ZERO_BLOCKS = (2, 4, 6, 7) |
| 15 | SPD_BLOCK_SIZE = 64 |
| 16 | SPD_BLOCK_HEADER_FMT = '<HHHH' |
| 17 | SPD_BLOCK_HEADER_SIZE = calcsize(SPD_BLOCK_HEADER_FMT) |
| 18 | spd_block_header = namedtuple( |
| 19 | 'spd_block_header', 'Type, Length, Key, Reserved') |
| 20 | |
| 21 | def parseargs(): |
| 22 | parser = argparse.ArgumentParser(description='Inject SPDs into APCB binaries') |
| 23 | parser.add_argument( |
| 24 | 'apcb_in', |
| 25 | type=str, |
| 26 | help='APCB input file') |
| 27 | parser.add_argument( |
| 28 | 'apcb_out', |
| 29 | type=str, |
| 30 | help='APCB output file') |
| 31 | parser.add_argument( |
| 32 | '--spd_sources', |
| 33 | nargs='+', |
| 34 | help='List of SPD sources') |
| 35 | return parser.parse_args() |
| 36 | |
| 37 | |
| 38 | # Calculate checksum of APCB binary |
| 39 | def chksum(data): |
| 40 | sum = 0 |
| 41 | for i, v in enumerate(data): |
| 42 | if i == APCB_CHECKSUM_OFFSET: continue |
| 43 | sum = (sum + v) & 0xff |
| 44 | return (0x100 - sum) & 0xff |
| 45 | |
| 46 | |
| 47 | # Inject bytes into binary blob by overwriting |
| 48 | def inject(orig, insert, offset): |
| 49 | return b''.join([orig[:offset], insert, orig[offset + len(insert):]]) |
| 50 | |
| 51 | |
| 52 | def main(): |
| 53 | args = parseargs() |
| 54 | |
| 55 | # Load input APCB |
| 56 | print(f'Reading input APCB from {args.apcb_in}') |
| 57 | with open(args.apcb_in, 'rb') as f: |
| 58 | apcb = f.read() |
| 59 | assert chksum(apcb) == apcb[APCB_CHECKSUM_OFFSET], 'Initial checksum is invalid' |
| 60 | orig_apcb_len = len(apcb) |
| 61 | |
| 62 | # Load SPDs |
| 63 | print(f'Using SPD Sources = {", ".join(args.spd_sources)}') |
| 64 | spds = [] |
| 65 | for spd_source in args.spd_sources: |
| 66 | with open(spd_source, 'rb') as f: |
| 67 | spd_data = bytes.fromhex(re.sub(r'\s+', '', f.read().decode())) |
| 68 | assert len(spd_data) == SPD_SIZE, f'{spd_source} is not {SPD_SIZE} bytes' |
| 69 | # Verify ZERO_BLOCKS are zero |
| 70 | for b in ZERO_BLOCKS: |
| 71 | assert all(v==0 for v in spd_data[b*SPD_BLOCK_SIZE:(b+1)*SPD_BLOCK_SIZE]), f'SPD block #{b} is not zero' |
| 72 | spds.append(spd_data) |
| 73 | assert len(spds) > 0, "No SPDs provided" |
| 74 | |
| 75 | # Inject SPDs into APCB |
| 76 | apcb_offset = 0 |
| 77 | spd_idx = 0 |
| 78 | while True: |
| 79 | apcb_offset = apcb.find(SPD_ENTRY_MAGIC, apcb_offset) |
| 80 | if apcb_offset < 0: |
| 81 | print(f'No more SPD entries found') |
| 82 | assert spd_idx >= len(spds), f'Not enough SPD entries in APCB. Need {len(spds)}, found {spd_idx}' |
| 83 | break |
| 84 | |
| 85 | if spd_idx < len(spds): |
| 86 | print(f'Injecting SPD instance {spd_idx}') |
| 87 | spd = spds[spd_idx] |
| 88 | else: |
| 89 | print(f'Injecting empty SPD for instance {spd_idx}') |
| 90 | spd = EMPTY_SPD |
| 91 | |
| 92 | # Inject SPD blocks |
| 93 | for b in range(int(SPD_SIZE/SPD_BLOCK_SIZE)): |
| 94 | if b in ZERO_BLOCKS: continue |
| 95 | header_data = apcb[apcb_offset:apcb_offset + SPD_BLOCK_HEADER_SIZE] |
| 96 | header = spd_block_header._make(unpack(SPD_BLOCK_HEADER_FMT, header_data)) |
| 97 | socket = (header.Key >> 12) & 0xF |
| 98 | channel = (header.Key >> 8) & 0xF |
| 99 | dimm = (header.Key >> 4) & 0xF |
| 100 | block_id = (header.Key >> 0) & 0xF |
| 101 | |
| 102 | assert header.Type == 2 |
| 103 | assert header.Length == SPD_BLOCK_HEADER_SIZE + SPD_BLOCK_SIZE |
| 104 | assert socket == 0 |
| 105 | assert channel == 0 |
| 106 | assert block_id == b |
| 107 | assert dimm == 0 |
| 108 | assert header.Reserved == 0 |
| 109 | |
| 110 | spd_block = spd[b*SPD_BLOCK_SIZE:(b+1)*SPD_BLOCK_SIZE] |
| 111 | apcb_offset += SPD_BLOCK_HEADER_SIZE |
| 112 | apcb = inject(apcb, spd_block, apcb_offset) |
| 113 | apcb_offset += SPD_BLOCK_SIZE |
| 114 | |
| 115 | spd_idx += 1 |
| 116 | |
| 117 | # Fix APCB checksum |
| 118 | print(f'Fixing APCB checksum') |
| 119 | apcb = inject(apcb, bytes([chksum(apcb)]), APCB_CHECKSUM_OFFSET) |
| 120 | assert chksum(apcb) == apcb[APCB_CHECKSUM_OFFSET], 'Final checksum is invalid' |
| 121 | assert orig_apcb_len == len(apcb), 'The size of the APCB changed.' |
| 122 | |
| 123 | # Write APCB to file |
| 124 | print(f'Writing {len(apcb)} byte APCB to {args.apcb_out}') |
| 125 | with open(args.apcb_out, 'wb') as f: |
| 126 | f.write(apcb) |
| 127 | |
| 128 | |
| 129 | if __name__ == "__main__": |
| 130 | main() |