import base64
import pathlib
import re

DUMP = "memory.dmp"

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

IEND_SIG = b"IEND\xAE\x42\x60\x82"  # конец PNG


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 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)
        if len(seg) > len(best):
            best = seg

    return best


def has_iend(b64_data):
    """Проверяем наличие IEND в декодированном бинаре."""
    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)}")

    # если первый кусок уже содержит IEND – редкость, но проверим
    if has_iend(merged):
        print("[+] IEND found in first block!")

    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

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

    # save 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 full
    if has_iend(merged):
        decoded = base64.b64decode(merged, validate=False)
        pathlib.Path("recovered.png").write_bytes(decoded)
        print("[+] PNG saved → recovered.png")
    else:
        print("[-] IEND not found — PNG incomplete.")


if __name__ == "__main__":
    recover_png()