import base64
import pathlib

DUMP = "memory.dmp"

START_MARK = b"![accounts](data:image/png;base64,"
BASE64_CHARS = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
OVERLAP = 50
WINDOW_FWD = 300000

IEND_SIG = b"IEND\xAE\x42\x60\x82"  # PNG tail signature


def find_all(data, sub):
    start = 0
    while True:
        idx = data.find(sub, start)
        if idx == -1:
            return
        yield idx
        start = idx + 1


def extract_initial_chunk(data, start):
    out = []
    for b in data[start:]:
        if b in BASE64_CHARS:
            out.append(b)
        else:
            break
    return bytes(out)


def cut_at_iend(b64_chunk: bytes) -> bytes:
    """
    Если chunk уже содержит конец PNG — обрезаем строго по IEND.
    Иначе возвращаем chunk как есть.
    """
    try:
        raw = base64.b64decode(b64_chunk, validate=False)
    except Exception:
        return b64_chunk

    pos = raw.find(IEND_SIG)
    if pos == -1:
        return b64_chunk  # IEND в этом куске нет — это ok

    # Обрезаем сырой PNG до конца
    raw = raw[: pos + len(IEND_SIG)]

    # Перекодируем обратно в clean base64
    return base64.b64encode(raw)


def find_continuation(data, chunk):
    tail = chunk[-OVERLAP:][:15]
    best = b""

    for idx in find_all(data, tail):
        window = data[idx + OVERLAP: idx + OVERLAP + WINDOW_FWD]

        buf = []
        for b in window:
            if b in BASE64_CHARS:
                buf.append(b)
            else:
                break

        seg = bytes(buf)

        # 💡 ВАЖНО: здесь мы режем мусор.
        seg = cut_at_iend(seg)

        if len(seg) > len(best):
            best = seg

    return best


def has_iend(b64_data):
    try:
        raw = base64.b64decode(b64_data, validate=False)
    except Exception:
        return False

    return IEND_SIG in raw


def recover_png():
    data = pathlib.Path(DUMP).read_bytes()

    start = data.find(START_MARK)
    if start == -1:
        print("[-] No PNG markdown found.")
        return

    start_b64 = start + len(START_MARK)
    print(f"[+] Base64 starts at {start_b64}")

    merged = bytearray(extract_initial_chunk(data, start_b64))
    print(f"[+] First chunk length: {len(merged)}")

    while True:
        print(f"[+] Looking for continuation (len={len(merged)})…")
        cont = find_continuation(data, merged)

        if not cont:
            print("[!] No continuation found; PNG may be truncated.")
            break

        print(f"[+] Found continuation ({len(cont)} bytes)")
        merged.extend(cont)

        if has_iend(merged):
            print("[+] IEND detected — PNG is complete.")
            break
        else:
            merged = merged[:-10]

    # save base64
    pathlib.Path("recovered_base64.txt").write_bytes(merged)
    print("[+] saved recovered_base64.txt")

    # decode raw binary
    try:
        raw = base64.b64decode(merged, validate=False)
        pathlib.Path("recovered_raw.bin").write_bytes(raw)
        print("[+] saved recovered_raw.bin")
    except Exception as e:
        print("[-] raw decode failed:", e)

    # save PNG
    if has_iend(merged):
        pathlib.Path("recovered.png").write_bytes(raw)
        print("[+] PNG saved → recovered.png")
    else:
        print("[-] IEND not found — PNG incomplete.")


if __name__ == "__main__":
    recover_png()