Inside Bitcoin

501            if None in ret:
502                ret = {ctx:enc for ctx, enc in ret.items()
503                       if ctx is None or enc.size < ret[None].size}
504            if hole:
505                ret = {ctx:enc for ctx, enc in ret.items() if ctx is None or ctx == 0}
506            return ret, hole
507        res, _ = recurse(self._trie)
508        return res[0] if 0 in res else res[None]
509
510    @staticmethod
511    def _from_binnode(binnode: _BinNode) -> "ASMap":
512        """Construct an ASMap object from a _BinNode. Internal use only."""
513        def recurse(node: _BinNode, default: int) -> list:
514            if node.ins == _Instruction.RETURN:
515                return [node.arg1]
516            if node.ins == _Instruction.JUMP:
517                return [recurse(node.arg1, default), recurse(node.arg2, default)]
518            if node.ins == _Instruction.MATCH:
519                val = node.arg1
520                sub = recurse(node.arg2, default)
521                while val >= 2:
522                    bit = val & 1
523                    val >>= 1
524                    if bit:
525                        sub = [[default], sub]
526                    else:
527                        sub = [sub, [default]]
528                return sub
529            assert node.ins == _Instruction.DEFAULT
530            return recurse(node.arg2, node.arg1)
531        ret = ASMap()
532        if binnode.ins != _Instruction.END:
533            #pylint: disable=protected-access
534            ret._set_trie(recurse(binnode, 0))
535        return ret
536
537    def to_binary(self, fill: bool = False) -> bytes:
538        """
539        Convert this ASMap object to binary.
540
541        Argument:
542            fill: permit the resulting binary encoder to contain mappers for
543                  unassigned subnets in this ASMap object. Doing so may
544                  reduce the size of the encoding.
545        Returns:
546            A bytes object with the encoding of this ASMap object.
547        """
548        bits: list[int] = []
549
550        def recurse(node: _BinNode) -> None:
551            _CODER_INS.encode(node.ins.value, bits)
552            if node.ins == _Instruction.RETURN:
553                _CODER_ASN.encode(node.arg1, bits)
554            elif node.ins == _Instruction.JUMP:
555                _CODER_JUMP.encode(node.arg1.size, bits)
556                recurse(node.arg1)
557                recurse(node.arg2)
558            elif node.ins == _Instruction.DEFAULT:
559                _CODER_ASN.encode(node.arg1, bits)
560                recurse(node.arg2)
561            else:
562                assert node.ins == _Instruction.MATCH
563                _CODER_MATCH.encode(node.arg1, bits)
564                recurse(node.arg2)
565
566        binnode = self._to_binnode(fill)
567        if binnode.ins != _Instruction.END:
568            recurse(binnode)
569
570        val = 0
571        nbits = 0
572        ret = []
573        for bit in bits:
574            val += (bit << nbits)
575            nbits += 1
576            if nbits == 8:
577                ret.append(val)
578                val = 0
579                nbits = 0
580        if nbits:
581            ret.append(val)
582        return bytes(ret)
583
584    @staticmethod
585    def from_binary(bindata: bytes) -> Optional["ASMap"]:
586        """Decode an ASMap object from the provided binary encoding."""
587
588        bits: list[int] = []
589        for byte in bindata:
590            bits.extend((byte >> i) & 1 for i in range(8))
591
592        def recurse(bitpos: int) -> tuple[_BinNode, int]:
593            insval, bitpos = _CODER_INS.decode(bits, bitpos)
594            ins = _Instruction(insval)
595            if ins == _Instruction.RETURN:
596                asn, bitpos = _CODER_ASN.decode(bits, bitpos)
597                return _BinNode(ins, asn), bitpos
598            if ins == _Instruction.JUMP:
599                jump, bitpos = _CODER_JUMP.decode(bits, bitpos)
600                left, bitpos1 = recurse(bitpos)
601                if bitpos1 != bitpos + jump:
602                    raise ValueError("Inconsistent jump")
603                right, bitpos = recurse(bitpos1)
604                return _BinNode(ins, left, right), bitpos
605            if ins == _Instruction.MATCH:
606                match, bitpos = _CODER_MATCH.decode(bits, bitpos)
607                sub, bitpos = recurse(bitpos)
608                return _BinNode(ins, match, sub), bitpos
609            assert ins == _Instruction.DEFAULT
610            asn, bitpos = _CODER_ASN.decode(bits, bitpos)
611            sub, bitpos = recurse(bitpos)
612            return _BinNode(ins, asn, sub), bitpos
613
614        if len(bits) == 0:
615            binnode = _BinNode(_Instruction.END)
616        else:
617            try:
618                binnode, bitpos = recurse(0)
619            except (ValueError, IndexError):
620                return None
621            if bitpos < len(bits) - 7:
622                return None
623            if not all(bit == 0 for bit in bits[bitpos:]):
624                return None
625
626        return ASMap._from_binnode(binnode)
627
628    def __lt__(self, other: "ASMap") -> bool:
629        return self._trie < other._trie
630
631    def __eq__(self, other: object) -> bool:
632        if isinstance(other, ASMap):
633            return self._trie == other._trie
634        return False
635
636    def extends(self, req: "ASMap") -> bool:
637        """Determine whether this matches req for all subranges where req is assigned."""
638        def recurse(actual: list, require: list) -> bool:
639            if len(require) == 1 and require[0] == 0:
640                return True
641            if len(require) == 1:
642                if len(actual) == 1:
643                    return bool(require[0] == actual[0])
644                return recurse(actual[0], require) and recurse(actual[1], require)
645            if len(actual) == 2:
646                return recurse(actual[0], require[0]) and recurse(actual[1], require[1])
647            return recurse(actual, require[0]) and recurse(actual, require[1])
648        assert isinstance(req, ASMap)
649        #pylint: disable=protected-access
650        return recurse(self._trie, req._trie)
651
652    def diff(self, other: "ASMap") -> list[ASNDiff]:
653        """Compute the diff from self to other."""
654        prefix: list[bool] = []
655        ret: list[ASNDiff] = []
656
657        def recurse(old_node: list, new_node: list):
658            if len(old_node) == 1 and len(new_node) == 1:
659                if old_node[0] != new_node[0]:
660                    ret.append((list(prefix), old_node[0], new_node[0]))
661            else:
662                old_left: list = old_node if len(old_node) == 1 else old_node[0]
663                old_right: list = old_node if len(old_node) == 1 else old_node[1]
664                new_left: list = new_node if len(new_node) == 1 else new_node[0]
665                new_right: list = new_node if len(new_node) == 1 else new_node[1]
666                prefix.append(False)
667                recurse(old_left, new_left)
668                prefix[-1] = True
669                recurse(old_right, new_right)
670                prefix.pop()
671        assert isinstance(other, ASMap)
672        #pylint: disable=protected-access
673        recurse(self._trie, other._trie)
674        return ret
675
676    def __copy__(self) -> "ASMap":
677        """Construct a copy of this ASMap object. Its state will not be shared."""
678        ret = ASMap()
679        #pylint: disable=protected-access
680        ret._set_trie(copy.deepcopy(self._trie))
681        return ret
682
683    def __deepcopy__(self, _) -> "ASMap":
684        # ASMap objects do not allow sharing of the _trie member, so we don't need the memoization.
685        return self.__copy__()
686
687
688class TestASMap(unittest.TestCase):
689    """Unit tests for this module."""
690
691    def test_ipv6_prefix_roundtrips(self) -> None:
692        """Test that random IPv6 network ranges roundtrip through prefix encoding."""
693        for _ in range(20):
694            net_bits = random.getrandbits(128)
695            for prefix_len in range(0, 129):
696                masked_bits = (net_bits >> (128 - prefix_len)) << (128 - prefix_len)
697                net = ipaddress.IPv6Network((masked_bits.to_bytes(16, 'big'), prefix_len))
698                prefix = net_to_prefix(net)
699                self.assertTrue(len(prefix) <= 128)
700                net2 = prefix_to_net(prefix)
701                self.assertEqual(net, net2)
702
703    def test_ipv4_prefix_roundtrips(self) -> None:
704        """Test that random IPv4 network ranges roundtrip through prefix encoding."""
705        for _ in range(100):
706            net_bits = random.getrandbits(32)
707            for prefix_len in range(0, 33):
708                masked_bits = (net_bits >> (32 - prefix_len)) << (32 - prefix_len)
709                net = ipaddress.IPv4Network((masked_bits.to_bytes(4, 'big'), prefix_len))
710                prefix = net_to_prefix(net)
711                self.assertTrue(32 <= len(prefix) <= 128)
712                net2 = prefix_to_net(prefix)
713                self.assertEqual(net, net2)
714
715    def test_asmap_roundtrips(self) -> None:
716        """Test case that verifies random ASMap objects roundtrip to/from entries/binary."""
717        # Iterate over the number of leaves the random test ASMap objects have.
718        for leaves in range(1, 20):
719            # Iterate over the number of bits in the AS numbers used.
720            for asnbits in range(0, 24):
721                # Iterate over the probability that leaves are unassigned.
722                for pct in range(101):
723                    # Construct a random ASMap object according to the above parameters.
724                    asmap = ASMap.from_random(num_leaves=leaves, max_asn=1 + (1 << asnbits),
725                                              unassigned_prob=0.01 * pct)
726                    # Run tests for to_entries and construction from those entries, both
727                    # for overlapping and non-overlapping ones.
728                    for overlapping in [False, True]:
729                        entries = asmap.to_entries(overlapping=overlapping, fill=False)
730                        random.shuffle(entries)
731                        asmap2 = ASMap(entries)
732                        assert asmap2 is not None
733                        self.assertEqual(asmap2, asmap)
734                        entries = asmap.to_entries(overlapping=overlapping, fill=True)
735                        random.shuffle(entries)
736                        asmap2 = ASMap(entries)
737                        assert asmap2 is not None
738                        self.assertTrue(asmap2.extends(asmap))
739
740                    # Run tests for to_binary and construction from binary.
741                    enc = asmap.to_binary(fill=False)
742                    asmap3 = ASMap.from_binary(enc)
743                    assert asmap3 is not None
744                    self.assertEqual(asmap3, asmap)
745                    enc = asmap.to_binary(fill=True)
746                    asmap3 = ASMap.from_binary(enc)
747                    assert asmap3 is not None
748                    self.assertTrue(asmap3.extends(asmap))
749
750    def test_patching(self) -> None: