import struct, sys, shutil, os


def align(x, a):
    return (x + a - 1) // a * a


def main():
    src, dst = sys.argv[1], sys.argv[2]
    size = int(sys.argv[3], 0) if len(sys.argv) > 3 else 0x300000
    secname = sys.argv[4] if len(sys.argv) > 4 else ".cjk"
    data = bytearray(open(src, "rb").read())

    e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
    coff = e_lfanew + 4
    nsec = struct.unpack_from("<H", data, coff + 2)[0]
    optsize = struct.unpack_from("<H", data, coff + 16)[0]
    opt = coff + 20
    imagebase = struct.unpack_from("<I", data, opt + 28)[0]
    sect_align = struct.unpack_from("<I", data, opt + 32)[0]
    file_align = struct.unpack_from("<I", data, opt + 36)[0]
    size_image = struct.unpack_from("<I", data, opt + 56)[0]
    so = opt + optsize

    secs = []
    for i in range(nsec):
        o = so + i * 40
        vs, va, rs, rp = struct.unpack_from("<IIII", data, o + 8)
        secs.append((va, vs, rp, rs, o))
    last_va, last_vs = secs[-1][0], secs[-1][1]
    new_va = align(last_va + last_vs, sect_align)
    new_raw = align(len(data), file_align)

    if so + (nsec + 1) * 40 > secs[0][2]:
        print("no room for new section header", file=sys.stderr)
        sys.exit(3)

    hdr = so + nsec * 40
    name = secname.encode("ascii")[:8].ljust(8, b"\x00")
    struct.pack_into("<8sIIII", data, hdr, name, size, new_va, align(size, file_align), new_raw)
    struct.pack_into("<IIHH", data, hdr + 36, 0xE0000020, 0, 0, 0)
    # characteristics: CNT_CODE|CNT_INITIALIZED_DATA|MEM_EXECUTE|MEM_READ|MEM_WRITE = 0xE0000020

    struct.pack_into("<H", data, coff + 2, nsec + 1)
    struct.pack_into("<I", data, opt + 56, align(new_va + size, sect_align))

    if len(data) < new_raw:
        data += b"\x00" * (new_raw - len(data))
    data += b"\x00" * align(size, file_align)

    open(dst, "wb").write(data)
    print("added %s VA=0x%x raw=0x%x size=0x%x -> %s (%d bytes)" % (secname, new_va, new_raw, size, dst, len(data)))
    print("section VA (absolute, imagebase 0x%x): 0x%x" % (imagebase, imagebase + new_va))


if __name__ == "__main__":
    main()