#!/usr/bin/env python3
"""Проверка сервера RustDesk по его протоколу, снаружи: check.py <host> <публичный ключ>.

1. UDP 21116 (hbbs): RegisterPeer -> RegisterPeerResponse - UDP доходит до контейнера и обратно.
2. TCP 21115 (hbbs): TestNatRequest -> TestNatResponse с нашим портом - NAT-тест работает, адрес клиента виден.
3. TCP 21116 (hbbs): PunchHoleRequest с нашим ключом -> ID_NOT_EXIST, с чужим -> LICENSE_MISMATCH:
   сервер работает именно с нашим ключом.
4. TCP 21117 (hbbr): два соединения RequestRelay с одним uuid сводятся, байты идут насквозь: relay работает.
   С чужим ключом hbbr соединение закрывает.
"""
import socket
import sys
import time
import uuid

# RendezvousMessage (hbb_common/protos/rendezvous.proto)
REGISTER_PEER, REGISTER_PEER_RESPONSE = 6, 7
PUNCH_HOLE_REQUEST, PUNCH_HOLE_RESPONSE = 8, 11
REQUEST_RELAY = 18
TEST_NAT_REQUEST, TEST_NAT_RESPONSE = 20, 21
FAILURES = {0: 'ID_NOT_EXIST', 2: 'OFFLINE', 3: 'LICENSE_MISMATCH', 4: 'LICENSE_OVERUSE'}
TIMEOUT = 5


def varint(value):
	out = bytearray()
	while True:
		byte, value = value & 0x7F, value >> 7
		out.append(byte | (0x80 if value else 0))
		if not value:
			return bytes(out)


def field(number, value):
	if isinstance(value, int):
		return varint(number << 3) + varint(value)
	value = value.encode() if isinstance(value, str) else value
	return varint(number << 3 | 2) + varint(len(value)) + value


def message(number, *fields):
	return field(number, b''.join(field(n, v) for n, v in fields))


def parse(data):
	"""{номер поля: значение} одного уровня: varint - int, length-delimited - bytes."""
	fields, i = {}, 0
	while i < len(data):
		key, i = read_varint(data, i)
		number, wire = key >> 3, key & 7
		if wire == 0:
			fields[number], i = read_varint(data, i)
		elif wire == 2:
			size, i = read_varint(data, i)
			fields[number], i = data[i:i + size], i + size
		else:
			raise ValueError(f'wire type {wire}')
	return fields


def read_varint(data, i):
	value = shift = 0
	while True:
		byte = data[i]
		value |= (byte & 0x7F) << shift
		i, shift = i + 1, shift + 7
		if not byte & 0x80:
			return value, i


def frame(payload):
	"""BytesCodec hbb_common: длина << 2 | (байт заголовка - 1), little endian."""
	n = len(payload)
	size = 1 if n <= 0x3F else 2 if n <= 0x3FFF else 3 if n <= 0x3FFFFF else 4
	return ((n << 2) | (size - 1)).to_bytes(size, 'little') + payload


def recv_exact(sock, n):
	data = b''
	while len(data) < n:
		chunk = sock.recv(n - len(data))
		if not chunk:
			raise ConnectionError('сервер закрыл соединение')
		data += chunk
	return data


def recv_frame(sock):
	first = recv_exact(sock, 1)
	size = (first[0] & 0x3) + 1
	head = first + recv_exact(sock, size - 1)
	return recv_exact(sock, int.from_bytes(head, 'little') >> 2)


def tcp(host, port, payload):
	with socket.create_connection((host, port), TIMEOUT) as sock:
		sock.settimeout(TIMEOUT)
		sock.sendall(frame(payload))
		return parse(recv_frame(sock)), sock.getsockname()[1]


def check_udp(host):
	with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
		sock.settimeout(TIMEOUT)
		sock.sendto(message(REGISTER_PEER, (1, f'probe{uuid.uuid4().hex[:8]}'), (2, 0)), (host, 21116))
		reply = parse(sock.recv(4096))
	assert REGISTER_PEER_RESPONSE in reply, reply
	return f'RegisterPeerResponse request_pk={parse(reply[REGISTER_PEER_RESPONSE]).get(2, 0)}'


def check_nat(host):
	reply, port = tcp(host, 21115, message(TEST_NAT_REQUEST, (1, 0)))
	seen = parse(reply[TEST_NAT_RESPONSE]).get(1)
	return f'TestNatResponse port={seen} (наш {port})'


def punch_hole(host, key):
	reply, _ = tcp(host, 21116, message(PUNCH_HOLE_REQUEST, (1, 'probe-nonexistent'), (3, key)))
	return FAILURES[parse(reply[PUNCH_HOLE_RESPONSE]).get(3, 0)]


def check_key(host, key):
	ours, other = punch_hole(host, key), punch_hole(host, 'wrong-key')
	assert (ours, other) == ('ID_NOT_EXIST', 'LICENSE_MISMATCH'), (ours, other)
	return f'наш ключ -> {ours}, чужой -> {other}'


def check_relay(host, key):
	session = uuid.uuid4().hex
	request = frame(message(REQUEST_RELAY, (2, session), (6, key)))
	with socket.create_connection((host, 21117), TIMEOUT) as a, socket.create_connection((host, 21117), TIMEOUT) as b:
		a.settimeout(TIMEOUT)
		b.settimeout(TIMEOUT)
		a.sendall(request)
		# второй должен прийти, когда первый уже ждет пару
		time.sleep(0.5)
		b.sendall(request)
		time.sleep(0.5)
		a.sendall(b'ping-from-a')
		got_b = recv_exact(b, len(b'ping-from-a'))
		b.sendall(b'pong-from-b')
		got_a = recv_exact(a, len(b'pong-from-b'))
	assert (got_a, got_b) == (b'pong-from-b', b'ping-from-a'), (got_a, got_b)
	with socket.create_connection((host, 21117), TIMEOUT) as c:
		c.settimeout(TIMEOUT)
		c.sendall(frame(message(REQUEST_RELAY, (2, uuid.uuid4().hex), (6, 'wrong-key'))))
		closed = c.recv(1) == b''
	assert closed, 'чужой ключ не отклонен'
	return 'два клиента сведены, байты прошли в обе стороны; чужой ключ - соединение закрыто'


def main(host, key):
	failed = 0
	for name, check in (
		('UDP 21116 hbbs', lambda: check_udp(host)),
		('TCP 21115 hbbs', lambda: check_nat(host)),
		('TCP 21116 hbbs', lambda: check_key(host, key)),
		('TCP 21117 hbbr', lambda: check_relay(host, key)),
	):
		try:
			print(f'OK   {name}: {check()}')
		except Exception as error:
			failed += 1
			print(f'FAIL {name}: {type(error).__name__}: {error}')
	return 1 if failed else 0


if __name__ == '__main__':
	sys.exit(main(*sys.argv[1:]))