#!/usr/bin/env python3
"""STUN Binding: проверяет UDP/DNS, не включает микрофон и не передаёт аудио.

Адрес XOR-MAPPED в выводе принадлежит машине запуска. Не публикуйте личный
сетевой адрес без очистки. --dry-run не разрешает DNS и не открывает сокет.
"""
import argparse
import ipaddress
import json
import secrets
import socket
import struct
import time

COOKIE = 0x2112A442


def request(transaction):
    # RFC 8489: 20-байтный заголовок Binding Request без атрибутов.
    if len(transaction) != 12:
        raise ValueError('Transaction ID должен занимать 12 байт')
    return struct.pack('!HHI', 0x0001, 0, COOKIE) + transaction


def decode(data, transaction):
    # Не считаем успехом короткий ответ, чужой transaction или неверную длину.
    if len(data) < 20:
        raise ValueError('Обрезан STUN header')
    kind, size, cookie = struct.unpack('!HHI', data[:8])
    if cookie != COOKIE or data[8:20] != transaction or len(data) != 20 + size or size % 4:
        raise ValueError('Неверные cookie, transaction или длина STUN')
    if kind not in (0x0101, 0x0111):
        raise ValueError('Это не Binding success/error response')
    result = {'response_type': hex(kind), 'success': kind == 0x0101}
    offset = 20
    while offset < len(data):
        if offset + 4 > len(data):
            raise ValueError('Обрезан header атрибута')
        attribute, length = struct.unpack('!HH', data[offset:offset + 4])
        end = offset + 4 + length
        padded = offset + 4 + ((length + 3) // 4) * 4
        if padded > len(data):
            raise ValueError('Обрезан STUN attribute')
        value = data[offset + 4:end]
        if attribute in (0x0020, 0x0001):
            if length < 4:
                raise ValueError('Обрезан MAPPED-ADDRESS')
            family = value[1]
            if family not in (1, 2) or length != (8 if family == 1 else 20):
                raise ValueError('Неверные family или длина MAPPED-ADDRESS')
            port = struct.unpack('!H', value[2:4])[0]
            raw = value[4:]
            if attribute == 0x0020:
                # XOR применяется к порту и IP, а не ко всему пакету.
                port ^= COOKIE >> 16
                mask = struct.pack('!I', COOKIE) + transaction
                raw = bytes(a ^ b for a, b in zip(raw, mask))
            result['xor_mapped' if attribute == 0x0020 else 'mapped'] = {
                'address': str(ipaddress.ip_address(raw)), 'port': port}
        elif attribute == 0x8022:
            result['software'] = value.decode('utf-8', 'replace')[:120]
        elif attribute == 0x0009:
            if length < 4:
                raise ValueError('Обрезан ERROR-CODE')
            result['error_code'] = (value[2] & 7) * 100 + value[3]
            result['reason'] = value[4:].decode('utf-8', 'replace')[:160]
        offset = padded
    if result['success'] and not any(k in result for k in ('mapped', 'xor_mapped')):
        raise ValueError('Binding success не содержит mapped address')
    return result


def probe(host, port=3478, family='4'):
    address_family = socket.AF_INET if family == '4' else socket.AF_INET6
    # Проверяем семейства раздельно. IPv6 DNS failure не означает IPv4 outage.
    addresses = socket.getaddrinfo(host, port, address_family, socket.SOCK_DGRAM)
    address = addresses[0][4]
    transaction = secrets.token_bytes(12)
    packet = request(transaction)
    started = time.monotonic()
    with socket.socket(address_family, socket.SOCK_DGRAM) as sock:
        sock.settimeout(2)
        # Connected UDP фильтрует ответы от постороннего отправителя; это не TCP connect.
        sock.connect(address)
        for attempt in range(1, 4):
            sock.send(packet)
            try:
                data = sock.recv(4096)
            except socket.timeout:
                if attempt == 3:
                    raise TimeoutError('Нет STUN ответа после трёх попыток')
                continue
            return {'host': host, 'family': family, 'server': list(address),
                    'request_bytes': len(packet), 'response_bytes': len(data),
                    'attempt': attempt, 'elapsed_ms': round((time.monotonic() - started) * 1000, 1),
                    **decode(data, transaction)}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('host')
    parser.add_argument('--port', type=int, default=3478)
    parser.add_argument('--family', choices=('4', '6'), default='4')
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()
    if not 1 <= args.port <= 65535:
        parser.error('Порт должен быть от 1 до 65535')
    if args.dry_run:
        print(json.dumps({'host': args.host, 'port': args.port, 'family': args.family,
                          'transport': 'UDP', 'request_bytes': 20, 'network': False,
                          'audio': False}, ensure_ascii=False))
        return
    try:
        result = probe(args.host, args.port, args.family)
    except socket.gaierror as exc:
        print(json.dumps({'stage': 'dns', 'family': args.family, 'error': str(exc)}, ensure_ascii=False))
        raise SystemExit(2)
    except (OSError, ValueError) as exc:
        print(json.dumps({'stage': 'udp/stun', 'family': args.family, 'error': str(exc)}, ensure_ascii=False))
        raise SystemExit(3)
    print(json.dumps(result, ensure_ascii=False))
    if not result['success']:
        raise SystemExit(4)


if __name__ == '__main__':
    main()