util/apcb/apcb_v3a_edit.py

   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()