hololake-system-architecture/product-source/hololake-platform/guanghu-os/scripts/qemu-native-net-peer.py

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


GUEST_MAC = bytes.fromhex("525400267198")
PEER_MAC = bytes.fromhex("525400123401")
GUEST_IP = socket.inet_aton("172.16.0.6")
GATEWAY_IP = socket.inet_aton("172.16.0.1")
RELAY_IP = socket.inet_aton("43.153.193.169")
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 arp_reply(request: bytes) -> bytes:
    assert request[12:14] == b"\x08\x06"
    assert request[20:22] == b"\x00\x01"
    sender_mac = request[22:28]
    sender_ip = request[28:32]
    target_ip = request[38:42]
    assert sender_mac == GUEST_MAC
    assert sender_ip == GUEST_IP
    assert target_ip == GATEWAY_IP
    return (
        sender_mac
        + PEER_MAC
        + b"\x08\x06"
        + b"\x00\x01\x08\x00\x06\x04\x00\x02"
        + PEER_MAC
        + GATEWAY_IP
        + sender_mac
        + sender_ip
    )


def verified_request(frame: bytes, sequence: int, magic: bytes) -> None:
    assert frame[0:6] == PEER_MAC
    assert frame[6:12] == GUEST_MAC
    assert frame[12:14] == b"\x08\x00"
    assert frame[26:30] == GUEST_IP
    assert frame[30:34] == RELAY_IP
    assert checksum(frame[14:34]) == 0
    assert frame[34] == 8
    assert frame[35] == 0
    assert frame[38:40] == b"\x47\x48"
    assert int.from_bytes(frame[40:42], "big") == sequence
    assert frame[50:66] == magic
    assert frame[66:82] == magic
    assert checksum(frame[34:82]) == 0


def authenticated_reply(frame: bytes, magic: bytes) -> bytes:
    reply = bytearray(frame)
    reply[0:6] = GUEST_MAC
    reply[6:12] = PEER_MAC
    reply[26:30] = RELAY_IP
    reply[30:34] = GUEST_IP
    reply[24:26] = b"\0\0"
    reply[24:26] = struct.pack("!H", checksum(bytes(reply[14:34])))
    reply[34] = 0
    reply[36:38] = b"\0\0"
    reply[66:82] = NATIVE_ACK_MAGIC
    reply[36:38] = struct.pack("!H", checksum(bytes(reply[34:82])))
    return bytes(reply)


def main() -> None:
    global GUEST_IP, GATEWAY_IP, RELAY_IP
    parser = argparse.ArgumentParser()
    parser.add_argument("--listen-port", type=int, required=True)
    parser.add_argument("--qemu-port", type=int, required=True)
    parser.add_argument("--receipt", required=True)
    parser.add_argument("--resident", action="store_true")
    parser.add_argument("--login-only", action="store_true")
    parser.add_argument("--guest-ip", default="172.16.0.6")
    parser.add_argument("--peer-ip", default="172.16.0.1")
    parser.add_argument("--relay-ip", default="43.153.193.169")
    args = parser.parse_args()
    GUEST_IP = socket.inet_aton(args.guest_ip)
    GATEWAY_IP = socket.inet_aton(args.peer_ip)
    RELAY_IP = socket.inet_aton(args.relay_ip)

    peer = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    peer.bind(("127.0.0.1", args.listen_port))
    peer.settimeout(0.2)
    qemu = ("127.0.0.1", args.qemu_port)
    deadline = time.monotonic() + 15
    arp_verified = False
    login_count = 0
    resident_login_count = 0
    commit_verified = False
    branch_verified = False
    recovery_verified = False

    while time.monotonic() < deadline:
        try:
            frame = peer.recv(4096)
        except TimeoutError:
            continue
        if frame[12:14] == b"\x08\x06":
            peer.sendto(arp_reply(frame), qemu)
            arp_verified = True
            continue
        if frame[12:14] != b"\x08\x00":
            continue
        sequence = int.from_bytes(frame[40:42], "big")
        if sequence <= 3:
            magic = LOGIN_MAGIC
            login_count += 1
        elif sequence == 4:
            magic = COMMIT_MAGIC
            commit_verified = True
        elif sequence == 5:
            magic = BRANCH_MAGIC
            branch_verified = True
        elif 6 <= sequence <= 15:
            magic = LOGIN_MAGIC
            resident_login_count += 1
        elif sequence == 16:
            magic = RECOVERY_MAGIC
            recovery_verified = True
        else:
            raise AssertionError(f"unexpected native sequence {sequence}")
        verified_request(frame, sequence, magic)
        peer.sendto(authenticated_reply(frame, magic), qemu)

        terminal_sequence = 3 if args.login_only else (16 if args.resident else 5)
        if sequence != terminal_sequence:
            continue
        with open(args.receipt, "w", encoding="utf-8") as output:
            output.write(
                "arp_gateway_reply: VERIFIED\n"
                "handshake_direction: NATIVE_INITIATED_OUTBOUND_ICMP\n"
                "authenticated_relay_reply: VERIFIED\n"
                "icmp_login_request_sent: true\n"
                "icmp_login_reply_verified: true\n"
                f"icmp_login_reply_count: {login_count}\n"
                f"code_commit_reply_verified: {str(commit_verified).lower()}\n"
                f"branch_move_reply_verified: {str(branch_verified).lower()}\n"
                f"resident_login_reply_count: {resident_login_count}\n"
                f"recovery_reply_verified: {str(recovery_verified).lower()}\n"
                "login_magic: HLDP-GHOS-LOGIN!\n"
            )
        return
    raise SystemExit("timed out waiting for native outbound ICMP exchange")


if __name__ == "__main__":
    main()