import argparse
import csv
import dataclasses
import hashlib
import mmap
import pathlib
import shutil


@dataclasses.dataclass
class PathItem:
    offset: int
    orig: bytes
    changed: bytes
    comment: str
    ghidra_addr: str

    def apply(self, mm: mmap.mmap):
        assert len(self.orig) == len(self.changed), "Lengths mismatch"
        s = slice(self.offset, self.offset + len(self.orig))
        assert mm[s] == self.orig, "Original bytes not match"
        mm[s] = self.changed
        print(f"Applied {self.comment} (ghidra addr: 0x{self.ghidra_addr})")


@dataclasses.dataclass
class Instructions:
    md5_summ: bytes
    items: list[PathItem]

    def check_file_md5(self, mm: mmap.mmap):
        digest = hashlib.md5(mm).digest()
        assert digest == self.md5_summ


def get_bytes(s: str):
    if s.startswith("0x"):
        return bytes.fromhex(s[2:])
    else:
        return s.encode("ascii")


def load_instructions(csv_instr):
    with open(csv_instr) as f:
        reader = csv.DictReader(f)
        instructions = None
        for row in reader:
            if instructions is None:
                instructions = Instructions(bytes.fromhex(row['md5_summ']), [])
            patch_item = PathItem(
                int(row['offset'], 16),
                get_bytes(row['orig']),
                get_bytes(row['patched']),
                comment=row['comment'],
                ghidra_addr=row['ghidra_addr']
            )
            instructions.items.append(patch_item)
    if instructions is None:
        raise RuntimeError("Unable to create instructions")
    return instructions


def main(src: pathlib.Path, csv_instr: pathlib.Path, dst_suffix):
    instructions = load_instructions(csv_instr)
    dst = src.parent / f'{src.stem}_{dst_suffix}{src.suffix}'
    assert not dst.exists(), f"{dst} is exists"
    shutil.copyfile(src, dst)
    cnt = 0
    with dst.open("r+b") as f:
        with mmap.mmap(f.fileno(), 0) as mm:
            instructions.check_file_md5(mm)
            for item in instructions.items:
                item.apply(mm)
                cnt += 1
    print(f"Done, applied {cnt} changes, saved to {dst}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("src", type=pathlib.Path, help="bin file")
    parser.add_argument("csv_instr", type=pathlib.Path, help="CVS instructions for patching")
    parser.add_argument("dst_suffix", help="Suffix for dst file")
    args = parser.parse_args()
    main(**vars(args))