#!/usr/bin/env python3
"""
Usage:
  # 1) Global find/replace (same-length)
  python3 patch.py in.bin out.bin --fr FIND_HEX REPLACE_HEX

  # 2) Replace content BETWEEN markers (keep markers)
  python3 patch.py in.bin out.bin --between START_HEX END_HEX REPLACE_HEX [--once]

  # 3) In-place mode (patch input directly)
  python3 patch.py in.bin --in-place --fr FIND_HEX REPLACE_HEX
  python3 patch.py in.bin --in-place --between START_HEX END_HEX REPLACE_HEX [--once]

Notes:
- HEX may include spaces/newlines; case-insensitive.
- In --fr mode, FIND/REPLACE lengths must match.
- In --between mode, REPLACE_HEX repeats/truncates to fit exactly between markers.
- --in-place modifies the input directly; a .bak backup is made by default.
"""

from __future__ import annotations
import argparse
import mmap
import re
import shutil
import sys
from pathlib import Path
from binascii import unhexlify

HEX_CLEAN_RE = re.compile(r"[^0-9a-fA-F]+")

def clean_hex(h: str) -> str:
    return HEX_CLEAN_RE.sub("", h).lower()

def hex2bytes_strict(h: str) -> bytes:
    h = clean_hex(h)
    if not h or (len(h) % 2):
        raise ValueError(f"Invalid HEX (empty/odd): {h!r}")
    try:
        return unhexlify(h)
    except Exception:
        raise ValueError(f"Invalid HEX: {h!r}")

def find_all(mm: mmap.mmap, needle: bytes, start: int = 0):
    pos = start
    while True:
        idx = mm.find(needle, pos)
        if idx == -1:
            break
        yield idx
        pos = idx + len(needle)

def patch_fr(mm: mmap.mmap, find_b: bytes, repl_b: bytes) -> int:
    if len(find_b) != len(repl_b):
        raise ValueError("FIND and REPLACE must be the same length in --fr mode.")
    count = 0
    for idx in find_all(mm, find_b, 0):
        mm[idx : idx + len(repl_b)] = repl_b
        count += 1
    return count

def patch_between(mm: mmap.mmap, start_b: bytes, end_b: bytes, pat_b: bytes, once: bool) -> int:
    if not pat_b:
        raise ValueError("REPLACE_HEX cannot be empty.")
    matches = 0
    offset = 0
    mlen = len(mm)

    while True:
        pos_s = mm.find(start_b, offset)
        if pos_s == -1:
            break
        inner_start = pos_s + len(start_b)
        pos_e = mm.find(end_b, inner_start)
        if pos_e == -1:
            break

        inner_len = pos_e - inner_start
        if inner_len > 0:
            rep = (pat_b * ((inner_len // len(pat_b)) + 2))[:inner_len]
            mm[inner_start : inner_start + inner_len] = rep
            matches += 1

        offset = pos_e + len(end_b)
        if once or offset >= mlen or offset <= pos_s:
            break

    return matches

def main(argv=None) -> int:
    p = argparse.ArgumentParser(description="Binary patcher for same-length or between-markers edits.")
    p.add_argument("inp", help="Input file")
    p.add_argument("out", nargs="?", help="Output file (omit with --in-place)")

    mode = p.add_mutually_exclusive_group(required=True)
    mode.add_argument("--fr", nargs=2, metavar=("FIND_HEX", "REPLACE_HEX"))
    mode.add_argument("--between", nargs=3, metavar=("START_HEX", "END_HEX", "REPLACE_HEX"))

    p.add_argument("--once", action="store_true", help="Patch only first between-block")
    p.add_argument("--in-place", action="store_true", help="Modify input file directly (no output file needed)")
    p.add_argument("--no-backup", action="store_true", help="Skip automatic .bak backup when using --in-place")

    args = p.parse_args(argv)

    in_path = Path(args.inp)
    if not in_path.is_file():
        print(f"Cannot read {in_path}", file=sys.stderr)
        return 1

    if args.in_place:
        out_path = in_path
        if not args.no_backup:
            bak_path = in_path.with_suffix(in_path.suffix + ".bak")
            shutil.copyfile(in_path, bak_path)
            print(f"Backup created: {bak_path}")
    else:
        if not args.out:
            print("Output file required unless using --in-place.", file=sys.stderr)
            return 1
        out_path = Path(args.out)
        shutil.copyfile(in_path, out_path)

    try:
        with out_path.open("r+b") as f:
            with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_WRITE) as mm:
                if args.fr:
                    find_hex, repl_hex = args.fr
                    find_b = hex2bytes_strict(find_hex)
                    repl_b = hex2bytes_strict(repl_hex)
                    count = patch_fr(mm, find_b, repl_b)
                    print(f"Replacements: {count}")
                elif args.between:
                    start_hex, end_hex, pat_hex = args.between
                    start_b = hex2bytes_strict(start_hex)
                    end_b = hex2bytes_strict(end_hex)
                    pat_b = hex2bytes_strict(pat_hex)
                    matches = patch_between(mm, start_b, end_b, pat_b, args.once)
                    print(f"Marker blocks patched: {matches}")
                mm.flush()
    except Exception as e:
        print(f"Patching failed: {e}", file=sys.stderr)
        return 1

    print(f"Wrote: {out_path}")
    return 0

if __name__ == "__main__":
    raise SystemExit(main())