"""Verify venue DNS over UDP/TCP; --docker tests an isolated server without WAN."""

import argparse
from pathlib import Path
import os
import socket
import struct
import subprocess
import time
import uuid


DOMAINS = ['evening-detective.crabs-games.art', 'evening-detective-api.crabs-games.art']
IMAGE = 'coredns/coredns:1.14.7'


def query(server, domain, kind=1, tcp=False):
    identifier = os.urandom(2)
    question = b''.join(bytes([len(label)]) + label.encode() for label in domain.split('.')) + b'\0'
    packet = identifier + struct.pack('!5H', 0x0100, 1, 0, 0, 0) + question + struct.pack('!2H', kind, 1)
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM if tcp else socket.SOCK_DGRAM) as connection:
        connection.settimeout(3)
        connection.connect((server, 53))
        if tcp:
            connection.sendall(struct.pack('!H', len(packet)) + packet)

            def receive(size):
                data = b''
                while len(data) < size:
                    chunk = connection.recv(size - len(data))
                    if not chunk:
                        raise RuntimeError('Incomplete DNS response')
                    data += chunk
                return data

            response = receive(struct.unpack('!H', receive(2))[0])
        else:
            connection.send(packet)
            response = connection.recv(4096)
    assert response[:2] == identifier, 'DNS response identifier mismatch'
    flags, questions, answers, _, _ = struct.unpack_from('!5H', response, 2)
    assert flags & 0x8000 and not flags & 0x0200, 'Invalid or truncated DNS response'

    def skip_name(offset):
        while response[offset]:
            if response[offset] & 0xc0 == 0xc0:
                return offset + 2
            offset += response[offset] + 1
        return offset + 1

    offset = 12
    for _ in range(questions):
        offset = skip_name(offset) + 4
    addresses = []
    for _ in range(answers):
        offset = skip_name(offset)
        record_type, _, _, length = struct.unpack_from('!HHIH', response, offset)
        offset += 10
        if record_type in (1, 28):
            addresses.append(socket.inet_ntop(socket.AF_INET if record_type == 1 else socket.AF_INET6,
                                              response[offset:offset + length]))
        offset += length
    return flags & 15, addresses


def verify(server, expected):
    for domain in DOMAINS:
        for tcp in (False, True):
            assert query(server, domain, tcp=tcp) == (0, [expected]), domain
            assert query(server, domain, kind=28, tcp=tcp) == (0, []), 'Unexpected external IPv6 answer: ' + domain
    print('PASS: both venue domains resolve locally over UDP/TCP, no external IPv6 answers')


def docker_test():
    root = Path(__file__).resolve().parents[2]
    network = 'venue-dns-test-' + uuid.uuid4().hex[:10]
    container = network + '-server'

    def run(*args):
        return subprocess.check_output(args, text=True, stderr=subprocess.PIPE).strip()

    try:
        run('docker', 'network', 'create', '--internal', network)
        run('docker', 'run', '-d', '--name', container, '--network', network,
            '-e', 'VENUE_DNS_IP=192.0.2.55', '-v', str(root / 'venue/Corefile') + ':/etc/coredns/Corefile:ro',
            IMAGE, '-conf', '/etc/coredns/Corefile')
        address = run('docker', 'inspect', '--format', '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}', container)
        for attempt in range(30):
            try:
                query(address, DOMAINS[0])
                break
            except (OSError, RuntimeError):
                time.sleep(0.1)
        else:
            raise AssertionError('DNS did not start')
        verify(address, '192.0.2.55')
        assert run('docker', 'network', 'inspect', '--format', '{{.Internal}}', network) == 'true'
        run('docker', 'restart', container)
        for attempt in range(30):
            try:
                query(address, DOMAINS[0])
                break
            except (OSError, RuntimeError):
                time.sleep(0.1)
        verify(address, '192.0.2.55')
        print('PASS: CoreDNS works without WAN before and after restart')
    finally:
        subprocess.run(['docker', 'rm', '-f', container], capture_output=True)
        subprocess.run(['docker', 'network', 'rm', network], capture_output=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--docker', action='store_true')
    parser.add_argument('--server', help='Laptop IPv4 address')
    args = parser.parse_args()
    if args.docker:
        docker_test()
    elif args.server:
        verify(args.server, args.server)
    else:
        parser.error('Use --docker or --server <laptop IPv4>')