252 lines
9.1 KiB
Python
Executable file
252 lines
9.1 KiB
Python
Executable file
#!/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()
|