1#!/usr/bin/env python
2# Fill in checksum/size of an option rom, and pad it to proper length.
3#
4# Copyright (C) 2009  Kevin O'Connor <kevin@koconnor.net>
5#
6# This file may be distributed under the terms of the GNU GPLv3 license.
7
8import sys, struct
9
10from python23compat import as_bytes
11
12def alignpos(pos, alignbytes):
13    mask = alignbytes - 1
14    return (pos + mask) & ~mask
15
16def checksum(data):
17    if (sys.version_info > (3, 0)):
18        cksum = sum(data)
19    else:
20        cksum = sum(map(ord, data))
21    return struct.pack('<B', (0x100 - cksum) & 0xff)
22
23def main():
24    inname = sys.argv[1]
25    outname = sys.argv[2]
26
27    # Read data in
28    f = open(inname, 'rb')
29    data = f.read()
30    f.close()
31    count = len(data)
32
33    # Pad to a 512 byte boundary
34    data += as_bytes("\0") * (alignpos(count, 512) - count)
35    count = len(data)
36
37    # Check if a pci header is present
38    pcidata = ord(data[24:25]) + (ord(data[25:26]) << 8)
39    if pcidata != 0:
40        blocks = struct.pack('<H', int(count/512))
41        data = data[:pcidata + 16] + blocks + data[pcidata + 18:]
42
43    # Fill in size field; clear checksum field
44    blocks = struct.pack('<B', int(count/512))
45    data = data[:2] + blocks + data[3:6] + as_bytes("\0") + data[7:]
46
47    # Checksum rom
48    data = data[:6] + checksum(data) + data[7:]
49
50    # Write new rom
51    f = open(outname, 'wb')
52    f.write(data)
53    f.close()
54
55if __name__ == '__main__':
56    main()
57