patch.py

#!/usr/bin/env python3
"""
MSI MAG B760 TOMAHAWK WIFI DDR4 (MS-7D96) BIOS v1.C0 stability patch

Targets the factory defaults stored in the BIOS image so that settings
applied on first boot (or after a CMOS clear) include stable voltages
for 4-DIMM DDR4 configurations.

Usage:
    python patch.py original.1C0 [patched.1C0]

If no output path is given, the patched file is written next to the
original with "_patched" appended to the name.
"""

import hashlib
import shutil
import struct
import sys
from pathlib import Path

# ── expected input ────────────────────────────────────────────────────────────

EXPECTED_SIZE = 33_554_432   # 32 MB SPI flash
EXPECTED_MD5  = "9afdca8872da20eef74c308d28835f3c"   # v1.C0 (E7D96IMS.1C0)

# ── patch definitions ─────────────────────────────────────────────────────────
#
# Each entry is:
#   variable name (matches NVAR ASCII name in the firmware)
#   expected variable size in bytes (safety check)
#   list of (offset_within_variable_data, new_bytes)
#
# Offsets match AMI IFR variable offsets extracted from the Setup PE32 driver.
# Values are little-endian millivolts for voltages, raw byte for C-state index.

PATCHES = {
    "Setup": {
        "expected_size": 5081,
        "changes": [
            # CPU SA (System Agent) voltage.  Low SA voltage starves the memory
            # controller on 4-DIMM configs, causing training failures and
            # instability well below rated DRAM speeds.
            (0x0CBF, b"\x46\x05"),   # 1350 mV  (factory default: 0 = auto)

            # CPU VDDQ -- voltage for the CPU-side DDR4 I/O termination.
            (0x0F10, b"\xE2\x04"),   # 1250 mV  (factory: auto)

            # DRAM VDDQ -- voltage for the DIMM-side DDR4 I/O.
            (0x0F98, b"\xE2\x04"),   # 1250 mV

            # Per-slot VDDQ overrides (A1, A2, B1, B2).
            # Without these the per-slot values stay at auto and can override
            # the global setting on some BIOS revisions.
            (0x0F9A, b"\xE2\x04"),   # slot A1
            (0x0F9C, b"\xE2\x04"),   # slot A2
            (0x0F9E, b"\xE2\x04"),   # slot B1
            (0x0FA0, b"\xE2\x04"),   # slot B2
        ],
    },
    "CpuSetup": {
        "expected_size": 961,
        "changes": [
            # Package C-state limit.  Allowing deep C-states (C6/C8/C10)
            # causes the SA and VDDQ rails to power-gate aggressively.
            # On already-marginal silicon (Intel Vmin Shift) this creates
            # voltage droop on wake that crashes the memory controller.
            # C3 max keeps the rails alive without disabling sleep entirely.
            (0x004B, b"\x02"),       # 0x02 = C3 limit  (factory: 0xFF = no limit)
        ],
    },
}

# ── NVAR parsing ──────────────────────────────────────────────────────────────

NVAR_VALID      = 0x80
NVAR_ASCII_NAME = 0x02
NVAR_GUID       = 0x04
NVAR_DATA_ONLY  = 0x08
NVAR_EXT_HEADER = 0x10

NVAR_REGION_START = 0x0100_0000
NVAR_REGION_SIZE  = 0x0010_0000   # 1 MB covers all NVAR stores in this image


def find_nvar_entries(image: bytearray, name: str) -> list[tuple[int, int, int]]:
    """Return list of (entry_offset, data_offset, data_len) for all valid NVAR
    entries with the given ASCII name, scanning for nested occurrences."""
    name_b = name.encode("ascii") + b"\x00"
    end    = NVAR_REGION_START + NVAR_REGION_SIZE
    found  = []
    pos    = NVAR_REGION_START

    while True:
        idx = image.find(b"NVAR", pos, end)
        if idx == -1:
            break
        pos = idx

        if pos + 10 >= end:
            break

        total = struct.unpack_from("<H", image, pos + 4)[0]
        if total < 12 or total > 0x1_0000:
            pos += 4
            continue

        flags = image[pos + 9]
        if not (flags & NVAR_VALID):
            pos += 4
            continue
        if flags & NVAR_DATA_ONLY:
            pos += 4
            continue
        if not (flags & NVAR_ASCII_NAME):
            pos += 4
            continue

        guid_inline = bool(flags & NVAR_GUID)
        name_off    = pos + 10 + (16 if guid_inline else 1)
        if name_off + len(name_b) >= end:
            pos += 4
            continue

        if image[name_off : name_off + len(name_b)] == name_b:
            data_off = name_off + len(name_b)
            data_len = (pos + total) - data_off
            if (flags & NVAR_EXT_HEADER) and data_len > 0:
                data_len -= image[pos + total - 1]
            found.append((pos, data_off, data_len))

        pos += 4

    return found


# ── main ──────────────────────────────────────────────────────────────────────

def patch(src: Path, dst: Path) -> None:
    data = bytearray(src.read_bytes())

    if len(data) != EXPECTED_SIZE:
        raise ValueError(
            f"unexpected file size {len(data):,} bytes (want {EXPECTED_SIZE:,}). "
            "Make sure this is the unmodified E7D96IMS.1C0 from MSI."
        )

    md5 = hashlib.md5(data).hexdigest()
    if md5 != EXPECTED_MD5:
        print(
            f"warning: MD5 {md5} does not match the expected v1.C0 checksum.\n"
            "The patch offsets were derived from v1.C0.  Applying them to a "
            "different version may brick the board.  Ctrl-C now to abort."
        )
        input("Press Enter to continue anyway...")

    total_writes = 0

    for var_name, spec in PATCHES.items():
        entries = find_nvar_entries(data, var_name)
        if not entries:
            print(f"  [!] {var_name}: no NVAR entries found -- skipping")
            continue

        for entry_off, data_off, data_len in entries:
            if data_len != spec["expected_size"]:
                print(
                    f"  [!] {var_name} at 0x{entry_off:08X}: "
                    f"size {data_len} != expected {spec['expected_size']} -- skipping"
                )
                continue

            for offset, value in spec["changes"]:
                if offset + len(value) > data_len:
                    print(f"  [!] {var_name} offset {offset:#x} out of range -- skipping")
                    continue
                data[data_off + offset : data_off + offset + len(value)] = value
                total_writes += 1

            print(
                f"  [+] {var_name:10s}  entry=0x{entry_off:08X}  "
                f"data=0x{data_off:08X}  size={data_len}"
            )

    dst.write_bytes(data)
    print(f"\nwrote {dst}  ({total_writes} values patched)")
    print(f"MD5: {hashlib.md5(data).hexdigest()}")


def main() -> None:
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    src = Path(sys.argv[1])
    if not src.exists():
        print(f"error: {src} not found")
        sys.exit(1)

    if len(sys.argv) >= 3:
        dst = Path(sys.argv[2])
    else:
        dst = src.with_name(src.stem + "_patched" + src.suffix)

    if dst.exists() and dst != src:
        print(f"output {dst} already exists, overwriting")

    print(f"input:  {src}")
    print(f"output: {dst}\n")
    patch(src, dst)


if __name__ == "__main__":
    main()