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