hololake-system-architecture/product-source/hololake-platform/guanghu-os/scripts/physical-native-icmp-peer.py

#!/usr/bin/env python3
import argparse
import pathlib
import socket
import struct
import time


LOGIN_MAGIC = b"HLDP-GHOS-LOGIN!"
COMMIT_MAGIC = b"HLDP-CODE-COMMIT"
BRANCH_MAGIC = b"HLDP-BRANCH-MOVE"
RECOVERY_MAGIC = b"HLDP-RECOVER-OS!"
NATIVE_ACK_MAGIC = b"HLDP-NATIVE-ACK!"


def checksum(payload: bytes) -> int:
    if len(payload) % 2:
        payload += b"\0"
    words = struct.unpack(f"!{len(payload) // 2}H", payload)
    total = sum(words)
    while total >> 16:
        total = (total & 0xFFFF) + (total >> 16)
    return (~total) & 0xFFFF


def request(sequence: int, magic: bytes) -> bytes:
    payload = b"\0" * 8 + magic + magic
    packet = struct.pack("!BBHHH", 8, 0, 0, 0x4748, sequence) + payload
    return packet[:2] + struct.pack("!H", checksum(packet)) + packet[4:]


def reply_sequence_index(packet: bytes) -> int:
    if packet and packet[0] >> 4 == 4:
        header_length = (packet[0] & 0x0F) * 4
        packet = packet[header_length:]
    if len(packet) < 8:
        raise RuntimeError("ICMP reply is too short")
    return int.from_bytes(packet[6:8], "big")


def verify_reply(packet: bytes, magic: bytes) -> int:
    if packet and packet[0] >> 4 == 4:
        header_length = (packet[0] & 0x0F) * 4
        packet = packet[header_length:]
    if len(packet) < 48 or packet[0] != 0 or packet[1] != 0:
        raise RuntimeError("unexpected ICMP reply shape")
    if packet[16:32] != magic:
        raise RuntimeError("ICMP reply does not carry the requested HLDP magic")
    if packet[32:48] != NATIVE_ACK_MAGIC:
        raise RuntimeError("ICMP reply does not carry the native ACK marker")
    if checksum(packet) != 0:
        raise RuntimeError("ICMP reply checksum failed")
    return int.from_bytes(packet[6:8], "big")


def exchange(
    peer: socket.socket,
    target: tuple[str, int],
    sequence: int,
    magic: bytes,
    deadline: float,
    retry_interval: float,
) -> None:
    while time.monotonic() < deadline:
        peer.sendto(request(sequence, magic), target)
        attempt_started = time.monotonic()
        try:
            packet, _ = peer.recvfrom(4096)
        except TimeoutError:
            time.sleep(max(0.0, retry_interval - (time.monotonic() - attempt_started)))
            continue
        try:
            verify_reply(packet, magic)
        except RuntimeError:
            time.sleep(max(0.0, retry_interval - (time.monotonic() - attempt_started)))
            continue
        return
    raise TimeoutError(f"no verified reply for sequence {sequence}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--target", required=True)
    parser.add_argument("--receipt", required=True)
    parser.add_argument("--timeout", type=float, default=90.0)
    parser.add_argument("--retry-interval", type=float, default=1.0)
    parser.add_argument("--raw-socket", action="store_true")
    parser.add_argument("--resident", action="store_true")
    parser.add_argument("--login-only", action="store_true")
    parser.add_argument("--resume-after-login", action="store_true")
    parser.add_argument("--resume-after-commit", action="store_true")
    parser.add_argument("--resume-after-branch", action="store_true")
    parser.add_argument("--resume-resident-count", type=int)
    parser.add_argument("--resident-pipeline", action="store_true")
    args = parser.parse_args()
    if args.retry_interval < 1.0:
        raise SystemExit("--retry-interval must be at least 1 second")
    if args.resume_resident_count is not None and not (
        0 <= args.resume_resident_count <= 10
    ):
        raise SystemExit("--resume-resident-count must be between 0 and 10")
    if sum(
        (
            args.resident,
            args.login_only,
            args.resume_after_login,
            args.resume_after_commit,
            args.resume_after_branch,
            args.resume_resident_count is not None,
            args.resident_pipeline,
        )
    ) > 1:
        raise SystemExit(
            "--resident, --login-only, --resume-after-login, and "
            "--resume-after-commit, --resume-after-branch, and "
            "--resume-resident-count, and --resident-pipeline are mutually exclusive"
        )

    socket_type = socket.SOCK_RAW if args.raw_socket else socket.SOCK_DGRAM
    peer = socket.socket(socket.AF_INET, socket_type, socket.IPPROTO_ICMP)
    peer.settimeout(0.2 if args.resident_pipeline else 1.0)
    target = (socket.gethostbyname(args.target), 0)
    deadline = time.monotonic() + args.timeout
    if args.resident_pipeline:
        pipeline = [
            LOGIN_MAGIC,
            LOGIN_MAGIC,
            LOGIN_MAGIC,
            COMMIT_MAGIC,
            BRANCH_MAGIC,
            *([LOGIN_MAGIC] * 10),
            RECOVERY_MAGIC,
        ]
        acknowledged: set[int] = set()
        while time.monotonic() < deadline and len(acknowledged) < len(pipeline):
            for pipeline_sequence, pipeline_magic in enumerate(pipeline, start=1):
                if time.monotonic() >= deadline:
                    break
                peer.sendto(request(pipeline_sequence, pipeline_magic), target)
                try:
                    packet, _ = peer.recvfrom(4096)
                    observed_sequence = reply_sequence_index(packet)
                    if not 1 <= observed_sequence <= len(pipeline):
                        raise RuntimeError("native ACK sequence is outside the pipeline")
                    reply_sequence = verify_reply(
                        packet,
                        pipeline[observed_sequence - 1],
                    )
                except (TimeoutError, RuntimeError, IndexError):
                    continue
                if reply_sequence == pipeline_sequence:
                    acknowledged.add(reply_sequence)
        if len(acknowledged) != len(pipeline):
            raise TimeoutError(
                "native pipeline incomplete; acknowledged sequences: "
                + ",".join(str(item) for item in sorted(acknowledged))
            )
        receipt = pathlib.Path(args.receipt)
        receipt.write_text(
            "schema: guanghu.physical-native-icmp-peer/v2\n"
            "status: PASS_100\n"
            f"target: {args.target}\n"
            "native_ack_marker: HLDP-NATIVE-ACK!\n"
            "acknowledged_sequences: 1-16\n"
            "login_reply_count: 3\n"
            "code_commit_reply_verified: true\n"
            "branch_move_reply_verified: true\n"
            "resident_login_reply_count: 10\n"
            "recovery_reply_verified: true\n",
            encoding="utf-8",
        )
        print(receipt.read_text(encoding="utf-8"), end="")
        return
    if args.resume_resident_count is not None:
        sequence = 6 + args.resume_resident_count
    elif args.resume_after_branch:
        sequence = 6
    elif args.resume_after_commit:
        sequence = 5
    elif args.resume_after_login:
        sequence = 4
    else:
        sequence = 1
    phases: list[tuple[str, bytes]] = []
    if not (
        args.resume_after_login
        or args.resume_after_commit
        or args.resume_after_branch
        or args.resume_resident_count is not None
    ):
        phases.extend(
            [
                ("login", LOGIN_MAGIC),
                ("login", LOGIN_MAGIC),
                ("login", LOGIN_MAGIC),
            ]
        )
    if (
        not args.login_only
        and not args.resume_after_commit
        and not args.resume_after_branch
        and args.resume_resident_count is None
    ):
        phases.extend(
            [
                ("code_commit", COMMIT_MAGIC),
                ("branch_move", BRANCH_MAGIC),
            ]
        )
    if args.resume_after_commit:
        phases.append(("branch_move", BRANCH_MAGIC))
    if args.resume_resident_count is not None:
        phases.extend(
            [("resident_login", LOGIN_MAGIC)]
            * (10 - args.resume_resident_count)
        )
        phases.append(("recovery", RECOVERY_MAGIC))
    elif (
        args.resident
        or args.resume_after_login
        or args.resume_after_commit
        or args.resume_after_branch
    ):
        phases.extend([("resident_login", LOGIN_MAGIC)] * 10)
        phases.append(("recovery", RECOVERY_MAGIC))

    completed: list[str] = []
    for phase, magic in phases:
        exchange(peer, target, sequence, magic, deadline, args.retry_interval)
        completed.append(phase)
        sequence += 1

    receipt = pathlib.Path(args.receipt)
    receipt.write_text(
        "schema: guanghu.physical-native-icmp-peer/v1\n"
        "status: PASS_100\n"
        f"target: {args.target}\n"
        f"login_reply_count: {completed.count('login')}\n"
        f"code_commit_reply_verified: {str('code_commit' in completed).lower()}\n"
        f"branch_move_reply_verified: {str('branch_move' in completed).lower()}\n"
        f"resident_login_resume_count: {args.resume_resident_count or 0}\n"
        f"resident_login_reply_count: {completed.count('resident_login')}\n"
        "resident_login_total_count: "
        f"{(args.resume_resident_count or 0) + completed.count('resident_login')}\n"
        f"recovery_reply_verified: {str('recovery' in completed).lower()}\n",
        encoding="utf-8",
    )
    print(receipt.read_text(encoding="utf-8"), end="")


if __name__ == "__main__":
    main()