#!/usr/bin/python3

# Copyright (C) Internet Systems Consortium, Inc. ("ISC")
#
# SPDX-License-Identifier: MPL-2.0
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0.  If a copy of the MPL was not distributed with this
# file, you can obtain one at https://mozilla.org/MPL/2.0/.
#
# See the COPYRIGHT file distributed with this work for additional
# information regarding copyright ownership.

from typing import List
from warnings import warn

from hypothesis.strategies import (
    binary,
    builds,
    composite,
    integers,
    just,
    nothing,
    permutations,
)

import dns.name
import dns.message
import dns.rdataclass
import dns.rdatatype

import isctest.name


@composite
def dns_names(
    draw,
    *,
    prefix: dns.name.Name = dns.name.empty,
    suffix: dns.name.Name = dns.name.root,
    min_labels: int = 1,
    max_labels: int = 128,
) -> dns.name.Name:
    """
    This is a hypothesis strategy to be used for generating DNS names with given `prefix`, `suffix`
    and with total number of labels specified by `min_labels` and `max labels`.

    For example, calling
    ```
    dns_names(
        prefix=dns.name.from_text("test"),
        suffix=dns.name.from_text("isc.org"),
        max_labels=6
    ).example()
    ```
    will result in names like `test.abc.isc.org.` or `test.abc.def.isc.org`.

    There is no attempt to make the distribution of the generated names uniform in any way.
    The strategy however minimizes towards shorter names with shorter labels.

    It can be used with to build compound strategies, like this one which generates random DNS queries.

    ```
    dns_queries = builds(
        dns.message.make_query,
        qname=dns_names(),
        rdtype=dns_rdatatypes,
        rdclass=dns_rdataclasses,
    )
    ```
    """

    prefix = prefix.relativize(dns.name.root)
    suffix = suffix.derelativize(dns.name.root)

    try:
        outer_name = prefix + suffix
        remaining_bytes = 255 - isctest.name.len_wire_uncompressed(outer_name)
        assert remaining_bytes >= 0
    except dns.name.NameTooLong:
        warn(
            "Maximal length name of name execeeded by prefix and suffix. Strategy won't generate any names.",
            RuntimeWarning,
        )
        return draw(nothing())

    minimum_number_of_labels_to_generate = max(0, min_labels - len(outer_name.labels))
    maximum_number_of_labels_to_generate = max_labels - len(outer_name.labels)
    if maximum_number_of_labels_to_generate < 0:
        warn(
            "Maximal number of labels execeeded by prefix and suffix. Strategy won't generate any names.",
            RuntimeWarning,
        )
        return draw(nothing())

    maximum_number_of_labels_to_generate = min(
        maximum_number_of_labels_to_generate, remaining_bytes // 2
    )
    if maximum_number_of_labels_to_generate < minimum_number_of_labels_to_generate:
        warn(
            f"Minimal number set to {minimum_number_of_labels_to_generate}, but in {remaining_bytes} bytes there is only space for maximum of {maximum_number_of_labels_to_generate} labels.",
            RuntimeWarning,
        )
        return draw(nothing())

    if remaining_bytes == 0 or maximum_number_of_labels_to_generate == 0:
        warn(
            f"Strategy will return only one name ({outer_name}) as it exactly matches byte or label length limit.",
            RuntimeWarning,
        )
        return draw(just(outer_name))

    chosen_number_of_labels_to_generate = draw(
        integers(
            minimum_number_of_labels_to_generate, maximum_number_of_labels_to_generate
        )
    )
    chosen_number_of_bytes_to_partion = draw(
        integers(2 * chosen_number_of_labels_to_generate, remaining_bytes)
    )
    chosen_lengths_of_labels = draw(
        _partition_bytes_to_labels(
            chosen_number_of_bytes_to_partion, chosen_number_of_labels_to_generate
        )
    )
    generated_labels = tuple(
        draw(binary(min_size=l - 1, max_size=l - 1)) for l in chosen_lengths_of_labels
    )

    return dns.name.Name(prefix.labels + generated_labels + suffix.labels)


RDATACLASS_MAX = RDATATYPE_MAX = 65535
try:
    dns_rdataclasses = builds(dns.rdataclass.RdataClass, integers(0, RDATACLASS_MAX))
    dns_rdatatypes = builds(dns.rdatatype.RdataType, integers(0, RDATATYPE_MAX))
except AttributeError:
    # In old dnspython versions, RDataTypes and RDataClasses are int and not enums.
    dns_rdataclasses = integers(0, RDATACLASS_MAX)  # type: ignore
    dns_rdatatypes = integers(0, RDATATYPE_MAX)  # type: ignore
dns_rdataclasses_without_meta = dns_rdataclasses.filter(dns.rdataclass.is_metaclass)

# NOTE: This should really be `dns_rdatatypes_without_meta = dns_rdatatypes_without_meta.filter(dns.rdatatype.is_metatype()`,
#       but hypothesis then complains about the filter being too strict, so it is done in a “constructive” way.
dns_rdatatypes_without_meta = integers(0, dns.rdatatype.OPT - 1) | integers(dns.rdatatype.OPT + 1, 127) | integers(256, RDATATYPE_MAX)  # type: ignore


@composite
def _partition_bytes_to_labels(
    draw, remaining_bytes: int, number_of_labels: int
) -> List[int]:
    two_bytes_reserved_for_label = 2

    # Reserve two bytes for each label
    partition = [two_bytes_reserved_for_label] * number_of_labels
    remaining_bytes -= two_bytes_reserved_for_label * number_of_labels

    assert remaining_bytes >= 0

    # Add a random number between 0 and the remainder to each partition
    for i in range(number_of_labels):
        added = draw(
            integers(0, min(remaining_bytes, 64 - two_bytes_reserved_for_label))
        )
        partition[i] += added
        remaining_bytes -= added

    # NOTE: Some of the remaining bytes will usually not be assigned to any label, but we don't care.

    return draw(permutations(partition))