1#!/usr/bin/env python3 2 3# Script for injecting SPDs into APCB_v3a binaries. 4 5import re 6import argparse 7from collections import namedtuple 8from struct import * 9 10APCB_CHECKSUM_OFFSET = 16 11SPD_ENTRY_MAGIC = bytes.fromhex('0200480000000000') 12SPD_SIZE = 512 13EMPTY_SPD = b'\x00' * SPD_SIZE 14ZERO_BLOCKS = (2, 4, 6, 7) 15SPD_BLOCK_SIZE = 64 16SPD_BLOCK_HEADER_FMT = '<HHHH' 17SPD_BLOCK_HEADER_SIZE = calcsize(SPD_BLOCK_HEADER_FMT) 18spd_block_header = namedtuple( 19 'spd_block_header', 'Type, Length, Key, Reserved') 20 21def 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 39def 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 48def inject(orig, insert, offset): 49 return b''.join([orig[:offset], insert, orig[offset + len(insert):]]) 50 51 52def 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 129if __name__ == "__main__": 130 main() 131