#!/usr/bin/env python3
"""Safely extend the v72 services.art inode by one block in Sony sparse system.img."""

from __future__ import annotations

import argparse
import hashlib
import os
import re
import shutil
import struct
import subprocess
from pathlib import Path


BLOCK_SIZE = 4096
SPARSE_MAGIC = 0xED26FF3A
CHUNK_RAW = 0xCAC1
CHUNK_DONT_CARE = 0xCAC3

PRE_RAW_SHA256 = "d6e7cb3c729fd33915dd093ff58bd672b80685d6c9688e532f1c8a8b8c37c5aa"
PRE_SPARSE_SHA256 = "3fa4aae88d9c31546cc92e5ba40c302e80d23e942ad49f3c162b702716c93260"
TRUNCATED_ART_SHA256 = "8827f8b5dc08001e5c491073548e00ebb0644366a3aa838d755954a3602da525"
FULL_ART_SHA256 = "600dd1ed1a3c8adc99c616b9037db8d222eb8b97b5db666b9adfc99d2d014175"

ART_PATH = "/system/framework/oat/arm/services.art"
ART_INODE = 2111
ART_MAIN_BLOCK = 141565
ART_MAIN_BLOCKS = 515
ART_EXTRA_LOGICAL_BLOCK = 515
ART_EXTRA_PHYSICAL_BLOCK = 408246
ART_TRUNCATED_SIZE = ART_MAIN_BLOCKS * BLOCK_SIZE
ART_FULL_SIZE = ART_TRUNCATED_SIZE + BLOCK_SIZE

# debugfs changes these metadata blocks when it allocates block 408246 and updates inode 2111.
ALLOCATION_METADATA_BLOCKS = (0, 1, 135, 393216)

EXPECTED_FILES = {
    "/system_ext/priv-app/SystemUI/SystemUI.apk":
        "5ddb2c94b96b18ac1d308270a1c0f8777e5b32de6cd8687aba516a0472756fc1",
    "/system_ext/priv-app/SystemUI/oat/arm/SystemUI.odex":
        "24a7a57bf671d852a03a608d175aee6960c9ca4c9968ef8b95a4428cb187bb9d",
    "/system_ext/priv-app/SystemUI/oat/arm/SystemUI.vdex":
        "26725319e171439a39fff84aac2a167d98fddcaa14d405981e8115cdbe7b5a49",
    "/system/framework/services.jar":
        "51978cfcbb44a2040117e996e3c8a97d00016e101429a6350a5d3e19e15cd514",
    "/system/framework/oat/arm/services.odex":
        "83cc66666d146c64f255afe61ca98e69a3950b69f14a37ace0d8364ec81c2121",
    "/system/framework/oat/arm/services.vdex":
        "16835f476f1e87293f77993cf078b4d478071acd4224c7275f9e2b59dc7d00a4",
    "/vendor/lib/hw/hwcomposer.msm8974.so":
        "16ea8a4436b8719ed3b8069b4a763383968b51940006385c915055d96b5b286d",
}


def run(*args: str, capture: bool = False, stdin: str | None = None) -> str:
    completed = subprocess.run(
        args,
        check=True,
        text=True,
        input=stdin,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.STDOUT if capture else None,
    )
    return completed.stdout or ""


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def dump_hash(raw: Path, image_path: str, target: Path) -> str:
    if target.exists():
        target.unlink()
    run("debugfs", "-R", f"dump {image_path} {target}", str(raw))
    result = sha256(target)
    target.unlink()
    return result


def sparse_segments(path: Path) -> list[tuple[int, int, int]]:
    result: list[tuple[int, int, int]] = []
    with path.open("rb") as sparse:
        header = sparse.read(28)
        (magic, major, minor, file_header_size, chunk_header_size, block_size,
         total_blocks, total_chunks, _checksum) = struct.unpack("<I4H4I", header)
        if (magic, major, minor, file_header_size, chunk_header_size, block_size) != (
            SPARSE_MAGIC, 1, 0, 28, 12, BLOCK_SIZE
        ):
            raise ValueError("unexpected sparse header")
        logical_block = 0
        for _ in range(total_chunks):
            chunk_start = sparse.tell()
            chunk_type, _reserved, chunk_blocks, total_size = struct.unpack(
                "<2H2I", sparse.read(12)
            )
            if chunk_type == CHUNK_RAW:
                if total_size != 12 + chunk_blocks * BLOCK_SIZE:
                    raise ValueError("invalid RAW chunk")
                result.append(
                    (
                        logical_block * BLOCK_SIZE,
                        (logical_block + chunk_blocks) * BLOCK_SIZE,
                        chunk_start + 12,
                    )
                )
            elif chunk_type == CHUNK_DONT_CARE:
                if total_size != 12:
                    raise ValueError("invalid DONT_CARE chunk")
            else:
                raise ValueError(f"Sony-incompatible sparse chunk 0x{chunk_type:04x}")
            logical_block += chunk_blocks
            sparse.seek(chunk_start + total_size)
        if logical_block != total_blocks:
            raise ValueError("sparse logical size mismatch")
    return result


def patch_range(
    sparse, segments: list[tuple[int, int, int]], raw_offset: int, data: bytes
) -> None:
    target_end = raw_offset + len(data)
    written = 0
    for segment_start, segment_end, payload_start in segments:
        start = max(raw_offset, segment_start)
        end = min(target_end, segment_end)
        if start >= end:
            continue
        sparse.seek(payload_start + start - segment_start)
        sparse.write(data[start - raw_offset:end - raw_offset])
        written += end - start
    if written != len(data):
        raise ValueError(f"range {raw_offset}:{target_end} is not fully RAW")


def copy_bytes(source, target, count: int) -> None:
    remaining = count
    while remaining:
        data = source.read(min(8 * 1024 * 1024, remaining))
        if not data:
            raise ValueError("unexpected end of sparse input")
        target.write(data)
        remaining -= len(data)


def inject_raw_block(sparse_path: Path, logical_target: int, data: bytes) -> None:
    if len(data) != BLOCK_SIZE:
        raise ValueError("injected sparse block must be exactly 4096 bytes")
    temporary = sparse_path.with_suffix(sparse_path.suffix + ".injecting")
    if temporary.exists():
        temporary.unlink()

    with sparse_path.open("rb") as source, temporary.open("w+b") as target:
        header = source.read(28)
        fields = list(struct.unpack("<I4H4I", header))
        magic, major, minor, file_header_size, chunk_header_size, block_size = fields[:6]
        total_blocks, total_chunks = fields[6], fields[7]
        if (magic, major, minor, file_header_size, chunk_header_size, block_size) != (
            SPARSE_MAGIC, 1, 0, 28, 12, BLOCK_SIZE
        ):
            raise ValueError("unexpected sparse header before injection")
        target.write(header)

        logical_block = 0
        output_chunks = 0
        injected = False
        for _ in range(total_chunks):
            chunk_header = source.read(12)
            chunk_type, reserved, chunk_blocks, total_size = struct.unpack(
                "<2H2I", chunk_header
            )
            chunk_end = logical_block + chunk_blocks
            contains_target = logical_block <= logical_target < chunk_end

            if not contains_target:
                target.write(chunk_header)
                copy_bytes(source, target, total_size - 12)
                output_chunks += 1
            else:
                if injected or chunk_type != CHUNK_DONT_CARE or total_size != 12:
                    raise ValueError("target block is not in one DONT_CARE chunk")
                before = logical_target - logical_block
                after = chunk_end - logical_target - 1
                if before:
                    target.write(struct.pack("<2H2I", CHUNK_DONT_CARE, reserved, before, 12))
                    output_chunks += 1
                target.write(
                    struct.pack("<2H2I", CHUNK_RAW, 0, 1, 12 + BLOCK_SIZE)
                )
                target.write(data)
                output_chunks += 1
                if after:
                    target.write(struct.pack("<2H2I", CHUNK_DONT_CARE, reserved, after, 12))
                    output_chunks += 1
                injected = True
            logical_block = chunk_end

        if logical_block != total_blocks or not injected:
            raise ValueError("sparse injection did not cover the target block")
        if source.read(1):
            raise ValueError("unexpected trailing sparse data")
        fields[7] = output_chunks
        target.seek(0)
        target.write(struct.pack("<I4H4I", *fields))
        target.flush()
        os.fsync(target.fileno())

    os.replace(temporary, sparse_path)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--raw", required=True, type=Path)
    parser.add_argument("--sparse", required=True, type=Path)
    parser.add_argument("--services-art", required=True, type=Path)
    parser.add_argument("--simg2img", required=True)
    args = parser.parse_args()

    if sha256(args.raw) != PRE_RAW_SHA256:
        raise ValueError("unexpected intermediate raw image")
    if sha256(args.sparse) != PRE_SPARSE_SHA256:
        raise ValueError("unexpected intermediate sparse image")
    if args.services_art.stat().st_size != ART_FULL_SIZE:
        raise ValueError("unexpected full services.art size")
    if sha256(args.services_art) != FULL_ART_SHA256:
        raise ValueError("unexpected full services.art hash")
    with args.services_art.open("rb") as source:
        prefix = source.read(ART_TRUNCATED_SIZE)
        extra_block = source.read(BLOCK_SIZE)
        if source.read(1):
            raise ValueError("unexpected services.art trailing data")
    if hashlib.sha256(prefix).hexdigest() != TRUNCATED_ART_SHA256:
        raise ValueError("services.art prefix does not match the intermediate image")

    stat_before = run(
        "debugfs", "-R", f"stat {ART_PATH}", str(args.raw), capture=True
    )
    expected_before = (
        f"Inode: {ART_INODE}",
        f"Size: {ART_TRUNCATED_SIZE}",
        f"(0-514):{ART_MAIN_BLOCK}-{ART_MAIN_BLOCK + ART_MAIN_BLOCKS - 1}",
    )
    if not all(marker in stat_before for marker in expected_before):
        raise ValueError(f"unexpected intermediate services.art inode:\n{stat_before}")
    if dump_hash(args.raw, ART_PATH, args.raw.with_name("verify-art-prefix")) != TRUNCATED_ART_SHA256:
        raise ValueError("intermediate services.art hash mismatch")
    free_check = run(
        "debugfs", "-R", f"testb {ART_EXTRA_PHYSICAL_BLOCK}", str(args.raw), capture=True
    )
    if "not in use" not in free_check:
        raise ValueError(f"expected extra block is not free: {free_check}")

    debugfs_commands = (
        f"fallocate {ART_PATH} {ART_EXTRA_LOGICAL_BLOCK} {ART_EXTRA_LOGICAL_BLOCK}\n"
        f"extent_open {ART_PATH}\n"
        f"goto_block {ART_EXTRA_LOGICAL_BLOCK}\n"
        f"set_bmap {ART_EXTRA_LOGICAL_BLOCK} {ART_EXTRA_PHYSICAL_BLOCK}\n"
        f"goto_block {ART_EXTRA_LOGICAL_BLOCK}\n"
        "current_node\n"
        "extent_close\n"
        f"set_inode_field {ART_PATH} size {ART_FULL_SIZE}\n"
        "quit\n"
    )
    run("debugfs", "-w", str(args.raw), capture=True, stdin=debugfs_commands)
    allocation_stat = run(
        "debugfs", "-R", f"stat {ART_PATH}", str(args.raw), capture=True
    )
    if (
        f"Size: {ART_FULL_SIZE}" not in allocation_stat
        or f"(515):{ART_EXTRA_PHYSICAL_BLOCK}" not in allocation_stat
        or f"(515[u]):{ART_EXTRA_PHYSICAL_BLOCK}" in allocation_stat
    ):
        raise ValueError(f"debugfs did not create the initialized extent:\n{allocation_stat}")

    with args.raw.open("r+b", buffering=0) as raw:
        raw.seek(ART_EXTRA_PHYSICAL_BLOCK * BLOCK_SIZE)
        raw.write(extra_block)

    segments = sparse_segments(args.sparse)
    with args.raw.open("rb") as raw, args.sparse.open("r+b", buffering=0) as sparse:
        for block in ALLOCATION_METADATA_BLOCKS:
            raw.seek(block * BLOCK_SIZE)
            patch_range(sparse, segments, block * BLOCK_SIZE, raw.read(BLOCK_SIZE))
    inject_raw_block(args.sparse, ART_EXTRA_PHYSICAL_BLOCK, extra_block)

    stat_after = run(
        "debugfs", "-R", f"stat {ART_PATH}", str(args.raw), capture=True
    )
    expected_after = (
        f"Inode: {ART_INODE}",
        f"Size: {ART_FULL_SIZE}",
        f"(0-514):{ART_MAIN_BLOCK}-{ART_MAIN_BLOCK + ART_MAIN_BLOCKS - 1}",
        f"(515):{ART_EXTRA_PHYSICAL_BLOCK}",
        'security.selinux (26) = "u:object_r:system_file:s0\\000"',
    )
    if not all(marker in stat_after for marker in expected_after):
        raise ValueError(f"unexpected final services.art inode:\n{stat_after}")
    if dump_hash(args.raw, ART_PATH, args.raw.with_name("verify-art-full")) != FULL_ART_SHA256:
        raise ValueError("final services.art hash mismatch")
    for image_path, expected_hash in EXPECTED_FILES.items():
        actual = dump_hash(args.raw, image_path, args.raw.with_name("verify-preserved"))
        if actual != expected_hash:
            raise ValueError(f"preserved file changed: {image_path}: {actual}")

    fsck = subprocess.run(
        ["e2fsck", "-fn", str(args.raw)],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    if fsck.returncode not in (0, 1):
        raise RuntimeError(fsck.stdout)

    roundtrip = args.raw.with_suffix(".roundtrip.raw")
    if roundtrip.exists():
        roundtrip.unlink()
    run(args.simg2img, str(args.sparse), str(roundtrip))
    if sha256(roundtrip) != sha256(args.raw):
        raise ValueError("Sony sparse roundtrip does not match final raw image")
    roundtrip.unlink()

    sparse_segments(args.sparse)
    print(f"services.art SHA-256: {FULL_ART_SHA256}")
    print(f"raw SHA-256: {sha256(args.raw)}")
    print(f"Sony sparse SHA-256: {sha256(args.sparse)}")
    print(f"services.art extent: (0-514):141565-142079, (515):{ART_EXTRA_PHYSICAL_BLOCK}")
    print("SELinux xattr preserved; all companion files exact")
    print("e2fsck clean; Sony sparse RAW/DONT_CARE only; sparse roundtrip exact")


if __name__ == "__main__":
    main()