WIP: copy from stepan/hypothesis 6250d2306a3c4910d175c4e720c17f84f772ed75

tests/dns/strategies.py

@@ -0,0 +1,165 @@
+#!/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
+
+# LATER: Move this file so it can be easily reused.
+
+
+@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 = 127,
+) -> 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 - len(outer_name.to_wire())
+        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 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 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 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 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
+dns_rdataclasses = builds(dns.rdataclass.RdataClass, integers(0, RDATACLASS_MAX))
+dns_rdataclasses_without_meta = dns_rdataclasses.filter(dns.rdataclass.is_metaclass)
+dns_rdatatypes = builds(dns.rdatatype.RdataType, integers(0, RDATATYPE_MAX))
+
+# 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))