lib.ips

  1import json
  2import random
  3from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Interface, IPv6Network
  4
  5import msgpack
  6from netaddr import EUI
  7
  8
  9class _AttrContainer:
 10    """Dynamic container for named attributes of one type (see the subclasses).
 11
 12    Serialises to ``{name: str(value)}`` (dict, JSON, msgpack) and back.
 13    """
 14    _item_type = object
 15    _type_name = "object"
 16
 17    def __setattr__(self, name, value):
 18        if not isinstance(value, self._item_type):
 19            raise TypeError(f"{name}: expected {self._type_name}, got {type(value).__name__}")
 20        super().__setattr__(name, value)
 21
 22    # ---------------- dict ----------------
 23
 24    def to_dict(self):
 25        return {k: str(v) for k, v in self.__dict__.items()}
 26
 27    @classmethod
 28    def from_dict(cls, d):
 29        obj = cls()
 30        for k, v in d.items():
 31            setattr(obj, k, cls._item_type(v))
 32        return obj
 33
 34    # ---------------- JSON ----------------
 35
 36    def to_json(self):
 37        return json.dumps(self.to_dict())
 38
 39    @classmethod
 40    def from_json(cls, s):
 41        return cls.from_dict(json.loads(s))
 42
 43    # ---------------- msgpack ----------------
 44
 45    def pack(self):
 46        return msgpack.packb(self.to_dict(), use_bin_type=True)
 47
 48    @classmethod
 49    def unpack(cls, blob):
 50        return cls.from_dict(msgpack.unpackb(blob, raw=False))
 51
 52
 53class IPv4Addresses(_AttrContainer):
 54    """Dynamic container for named IPv4Interface attributes, with no boilerplate.
 55
 56    Usage:
 57        ips = IPv4Addresses()
 58        ips.ip1 = IPv4Interface("192.168.1.1/24")
 59        ips.ip2 = IPv4Interface("192.168.1.2/24")
 60    """
 61    _item_type = IPv4Interface
 62    _type_name = "IPv4Interface"
 63
 64
 65class IPv4Networks(_AttrContainer):
 66    """Dynamic container for named IPv4Network attributes, with no boilerplate.
 67
 68    Usage:
 69        nets = IPv4Networks()
 70        nets.lan = IPv4Network("192.168.1.0/24")
 71        nets.mgmt = IPv4Network("10.0.0.0/8")
 72    """
 73    _item_type = IPv4Network
 74    _type_name = "IPv4Network"
 75
 76
 77class IPv6Addresses(_AttrContainer):
 78    """Dynamic container for named IPv6Interface attributes (the IPv6 twin of IPv4Addresses).
 79
 80    Usage:
 81        ips6 = IPv6Addresses()
 82        ips6.router = IPv6Interface("2001:db8:1::1/64")
 83    """
 84    _item_type = IPv6Interface
 85    _type_name = "IPv6Interface"
 86
 87
 88class IPv6Networks(_AttrContainer):
 89    """Dynamic container for named IPv6Network attributes (the IPv6 twin of IPv4Networks).
 90
 91    Usage:
 92        nets6 = IPv6Networks()
 93        nets6.lan = IPv6Network("2001:db8:1::/64")
 94    """
 95    _item_type = IPv6Network
 96    _type_name = "IPv6Network"
 97
 98
 99_PRIVATE_NETWORKS = [
100    IPv4Network("10.0.0.0/8"),
101    IPv4Network("172.16.0.0/12"),
102    IPv4Network("192.168.0.0/16"),
103]
104_PRIVATE_NETWORKS_V6 = [IPv6Network("fd00::/8")]  # RFC 4193 unique local addresses
105_GLOBAL_UNICAST_V6 = IPv6Network("2000::/3")  # default search space of random_ipv6networks
106_PRIVATE_BY_VERSION = {4: _PRIVATE_NETWORKS, 6: _PRIVATE_NETWORKS_V6}
107
108
109def _pick_one(mask, from_network, exclude, from_private_network):
110    """Pick a single random /mask network of the family of *from_network*.
111
112    Internal helper for random_ipv4networks / random_ipv6networks.
113    """
114    max_prefixlen = from_network.max_prefixlen
115    if not 0 <= mask <= max_prefixlen:
116        raise ValueError(f"mask /{mask} out of range for IPv{from_network.version} (0-{max_prefixlen})")
117    block_size = 1 << (max_prefixlen - mask)
118    net_cls = type(from_network)  # not ip_network((int, mask)): it would guess the family from the int
119    # networks of the other family never overlap: drop them
120    exclude = [ex for ex in exclude if ex.version == from_network.version]
121
122    def make(abs_idx):
123        return net_cls((abs_idx * block_size, mask))
124
125    def available(net):
126        return not any(net.overlaps(ex) for ex in exclude)
127
128    def indices_in(a, b):
129        first = (a + block_size - 1) // block_size
130        last = (b + 1) // block_size - 1
131        return range(first, last + 1) if first <= last else range(0, 0)
132
133    flo = int(from_network.network_address)
134    fhi = int(from_network.broadcast_address)
135    if from_private_network:
136        boxes = []
137        for p in _PRIVATE_BY_VERSION[from_network.version]:
138            lo = max(flo, int(p.network_address))
139            hi = min(fhi, int(p.broadcast_address))
140            if lo <= hi:
141                boxes.append((lo, hi))
142    else:
143        boxes = [(flo, fhi)]
144
145    # r.stop - r.start, not len(r): len() overflows for ranges longer than 2**63 (IPv6)
146    spans = [(r.start, r.stop - r.start) for lo, hi in boxes for r in [indices_in(lo, hi)] if r]
147    total = sum(count for _, count in spans)
148
149    if total == 0:
150        raise ValueError(f"No /{mask} network fits in the search space")
151
152    def pick_abs_idx():
153        if from_private_network and len(spans) > 1:
154            # Pick uniformly across private ranges (not weighted by range size)
155            start, count = random.choice(spans)
156            return start + random.randrange(count)
157        n = random.randrange(total)
158        for start, count in spans:
159            if n < count:
160                return start + n
161            n -= count
162        return spans[-1][0]  # unreachable: n < total guarantees a match above
163
164    if total <= 65536:
165        candidates = [make(i) for start, count in spans
166                      for i in range(start, start + count)
167                      if available(make(i))]
168        if not candidates:
169            raise ValueError(f"No available /{mask} network with the given constraints")
170        return random.choice(candidates)
171
172    for _ in range(1000):
173        candidate = make(pick_abs_idx())
174        if available(candidate):
175            return candidate
176
177    raise ValueError(f"No available /{mask} network found after 1000 attempts")
178
179
180def _random_networks(masks, from_network, exclude, from_private_network):
181    if isinstance(masks, int):
182        masks = [masks]
183
184    working_exclude = list(exclude) if exclude else []
185    result = []
186    for mask in masks:
187        net = _pick_one(mask, from_network, working_exclude, from_private_network)
188        result.append(net)
189        working_exclude.append(net)
190    return result
191
192
193def _check_version(func_name, network, version):
194    if getattr(network, 'version', None) != version:
195        raise ValueError(f"{func_name}: expected an IPv{version} network, got {type(network).__name__}")
196
197
198def random_ipv4networks(
199    masks,
200    from_network=IPv4Network("0.0.0.0/0"),
201    exclude=None,
202    from_private_network=False,
203):
204    """Return a list of disjoint random IPv4Networks.
205
206    `masks` is an int or a list of ints; one network is returned per mask,
207    all mutually disjoint and not overlapping any network in `exclude`.
208    Optionally restricted to RFC-1918 private ranges with `from_private_network=True`.
209
210    `from` being a reserved keyword, the parameter is named `from_network`.
211
212    Raises ValueError if any network cannot be allocated.
213    """
214    _check_version("random_ipv4networks", from_network, 4)
215    return _random_networks(masks, from_network, exclude, from_private_network)
216
217
218def random_ipv6networks(
219    masks,
220    from_network=None,
221    exclude=None,
222    from_private_network=False,
223):
224    """Return a list of disjoint random IPv6Networks (the IPv6 twin of random_ipv4networks).
225
226    `masks` is an int or a list of ints (prefix lengths, e.g. 64); one network is returned
227    per mask, all mutually disjoint and not overlapping any network in `exclude`.
228
229    `from_network=None` (default) means the global unicast space ``2000::/3``, or ``::/0``
230    when `from_private_network=True`: the search space is then the unique local range
231    ``fd00::/8`` (RFC 4193).  Use ``from_network=IPv6Network('2001:db8::/32')`` for
232    documentation prefixes.
233
234    Raises ValueError if any network cannot be allocated.
235    """
236    if from_network is None:
237        from_network = IPv6Network("::/0") if from_private_network else _GLOBAL_UNICAST_V6
238    _check_version("random_ipv6networks", from_network, 6)
239    return _random_networks(masks, from_network, exclude, from_private_network)
240
241
242def _host_offsets(network):
243    """Offsets of *network* usable as host addresses.
244
245    IPv4: the network and broadcast addresses are skipped up to /30; /31 (RFC 3021) and /32
246    use every address.  IPv6: offset 0 (the subnet-router anycast address, RFC 4291) is
247    skipped up to /126; /127 (RFC 6164) and /128 use every address.
248    """
249    total = network.num_addresses
250    if network.version == 4:
251        return range(1, total - 1) if network.prefixlen <= 30 else range(total)
252    return range(1, total) if network.prefixlen <= 126 else range(total)
253
254
255def _iface_cls(network):
256    return IPv4Interface if network.version == 4 else IPv6Interface
257
258
259def _random_addresses(network, n, exclude_ips, exclude_nets):
260    ex_ips = set(exclude_ips) if exclude_ips else set()
261    ex_nets = list(exclude_nets) if exclude_nets else []
262
263    base = int(network.network_address)
264    total = network.num_addresses
265    iface_cls = _iface_cls(network)
266
267    def make(i):
268        return iface_cls((base + i, network.prefixlen))
269
270    def available(ip):
271        return ip not in ex_ips and not any(ip in net for net in ex_nets)
272
273    host_range = _host_offsets(network)
274
275    # Small space: enumerate all candidates, then sample.
276    if total <= 65536:
277        candidates = [make(i) for i in host_range if available(make(i))]
278        if len(candidates) < n:
279            raise ValueError(
280                f"Not enough available addresses: need {n}, found {len(candidates)}"
281            )
282        return random.sample(candidates, n)
283
284    # Large space: pick one at a time, accumulating into a working exclude set.
285    working_ex = set(ex_ips)
286    result = []
287    for _ in range(n):
288        for _ in range(1000):
289            candidate = make(random.randrange(host_range.start, host_range.stop))
290            if candidate not in working_ex and not any(candidate in net for net in ex_nets):
291                result.append(candidate)
292                working_ex.add(candidate)
293                break
294        else:
295            raise ValueError(
296                f"Could not find enough available addresses (got {len(result)}/{n})"
297            )
298    return result
299
300
301def random_ipv4s(network, n=1, exclude_ips=None, exclude_nets=None):
302    """Return a list of n distinct random IPv4Interface within `network`,
303    excluding any address in `exclude_ips` or covered by any network in `exclude_nets`.
304
305    The network address and the broadcast address are never returned (prefix <= /30).
306    Note: `exclude_ips` is compared with the prefix, so list IPv4Interface values
307    (as returned by this function), not bare IPv4Address values.
308
309    Raises ValueError if fewer than n addresses are available.
310    """
311    _check_version("random_ipv4s", network, 4)
312    return _random_addresses(network, n, exclude_ips, exclude_nets)
313
314
315def random_ipv6s(network, n=1, exclude_ips=None, exclude_nets=None):
316    """Return a list of n distinct random IPv6Interface within `network` (the IPv6 twin
317    of random_ipv4s), excluding any address in `exclude_ips` or covered by any network
318    in `exclude_nets`.
319
320    Offset 0 (the subnet-router anycast address) is never returned up to /126; /127 and
321    /128 use every address.
322
323    Raises ValueError if fewer than n addresses are available.
324    """
325    _check_version("random_ipv6s", network, 6)
326    return _random_addresses(network, n, exclude_ips, exclude_nets)
327
328
329def _random_addresses_with_range(network, gap, n, exclude_ips, exclude_nets):
330    """Return a list of n + 2*k distinct random interfaces within `network`.
331
332    `gap` is either an int or a list of ints. ``k = 1`` when gap is an int,
333    ``k = len(gap)`` otherwise.
334
335    The returned list has the form:
336        [ip_min1, ip_max1, ip_min2, ip_max2, ..., ip_mink, ip_maxk, ip1, ..., ipn]
337
338    Guarantees:
339    - int(ip_max_i.ip) - int(ip_min_i.ip) == gap[i]  (or gap when gap is an int)
340    - ip_max_i < ip_min_{i+1}  (ranges are strictly ordered, non-overlapping)
341    - ip1 .. ipn are outside every range [ip_min_i, ip_max_i]
342
343    Raises ValueError if the constraints cannot be satisfied.
344    """
345    gaps = [gap] if isinstance(gap, int) else list(gap)
346    k = len(gaps)
347
348    ex_ips = set(exclude_ips) if exclude_ips else set()
349    ex_nets = list(exclude_nets) if exclude_nets else []
350
351    base = int(network.network_address)
352    total = network.num_addresses
353    prefixlen = network.prefixlen
354    iface_cls = _iface_cls(network)
355
356    def make(i):
357        return iface_cls((base + i, prefixlen))
358
359    def is_excluded(ip):
360        return ip in ex_ips or any(ip in net for net in ex_nets)
361
362    # Minimum space: k ranges placed back-to-back with 1-address gaps between them.
363    # Minimum last offset: sum(gaps) + k - 1; need total >= sum(gaps) + k.
364    min_needed = sum(gaps) + k
365    if total < min_needed:
366        raise ValueError(
367            f"Network {network} is too small to fit {k} range(s) with gaps {gaps}"
368        )
369
370    # suffix_sums[i] = sum(gaps[i:])
371    suffix_sums = [0] * (k + 1)
372    for i in range(k - 1, -1, -1):
373        suffix_sums[i] = suffix_sums[i + 1] + gaps[i]
374
375    def hi_for(i):
376        """Max start offset for range i that still leaves room for ranges i+1..k-1."""
377        return total - 1 - suffix_sums[i] - (k - 1 - i)
378
379    SMALL = 65536
380
381    def try_place_ranges():
382        """Try one left-to-right placement. Returns [(start, end), ...] or None."""
383        range_offsets = []
384        lo = 0
385        for i in range(k):
386            h = hi_for(i)
387            if h < lo:
388                return None
389            g = gaps[i]
390            space = h - lo + 1
391            if space <= SMALL:
392                candidates = [
393                    j
394                    for j in range(lo, h + 1)
395                    if not is_excluded(make(j)) and not is_excluded(make(j + g))
396                ]
397                if not candidates:
398                    return None
399                s = random.choice(candidates)
400            else:
401                s = None
402                for _ in range(1000):
403                    c = random.randint(lo, h)
404                    if not is_excluded(make(c)) and not is_excluded(make(c + g)):
405                        s = c
406                        break
407                if s is None:
408                    return None
409            range_offsets.append((s, s + g))
410            lo = s + g + 1
411        return range_offsets
412
413    range_offsets = None
414    for _ in range(1000):
415        range_offsets = try_place_ranges()
416        if range_offsets is not None:
417            break
418    if range_offsets is None:
419        raise ValueError(
420            f"Could not place {k} range(s) with gaps {gaps} with the given exclusions"
421        )
422
423    def in_any_range(offset):
424        return any(s <= offset <= e for s, e in range_offsets)
425
426    def extra_available(offset):
427        ip = make(offset)
428        return (
429            ip not in ex_ips
430            and not in_any_range(offset)
431            and not any(ip in net for net in ex_nets)
432        )
433
434    if total <= SMALL:
435        extra_candidates = [make(i) for i in range(total) if extra_available(i)]
436        if len(extra_candidates) < n:
437            raise ValueError(
438                f"Not enough addresses outside the range(s) for {n} extra IPs: "
439                f"found {len(extra_candidates)}"
440            )
441        extras = random.sample(extra_candidates, n)
442    else:
443        chosen = set()
444        extras = []
445        for _ in range(n):
446            for _ in range(1000):
447                offset = random.randrange(total)
448                candidate = make(offset)
449                if extra_available(offset) and candidate not in chosen:
450                    extras.append(candidate)
451                    chosen.add(candidate)
452                    break
453            else:
454                raise ValueError(
455                    f"Could not find enough extra addresses (got {len(extras)}/{n})"
456                )
457
458    result = []
459    for s, e in range_offsets:
460        result.append(make(s))
461        result.append(make(e))
462    result.extend(extras)
463    return result
464
465
466def random_ipv4s_with_range(network, gap, n=1, exclude_ips=None, exclude_nets=None):
467    """Return a list of n + 2*k distinct random IPv4Interface within `network`.
468
469    `gap` is either an int or a list of ints. ``k = 1`` when gap is an int,
470    ``k = len(gap)`` otherwise.
471
472    The returned list has the form:
473        [ip_min1, ip_max1, ip_min2, ip_max2, ..., ip_mink, ip_maxk, ip1, ..., ipn]
474
475    Guarantees:
476    - int(ip_max_i.ip) - int(ip_min_i.ip) == gap[i]  (or gap when gap is an int)
477    - ip_max_i < ip_min_{i+1}  (ranges are strictly ordered, non-overlapping)
478    - ip1 .. ipn are outside every range [ip_min_i, ip_max_i]
479
480    Unlike random_ipv4s, every offset of the network is eligible (including the first
481    and the last address).
482
483    Raises ValueError if the constraints cannot be satisfied.
484    """
485    _check_version("random_ipv4s_with_range", network, 4)
486    return _random_addresses_with_range(network, gap, n, exclude_ips, exclude_nets)
487
488
489def random_ipv6s_with_range(network, gap, n=1, exclude_ips=None, exclude_nets=None):
490    """IPv6 twin of random_ipv4s_with_range: n + 2*k distinct random IPv6Interface within
491    `network`, the first 2*k forming k ordered ranges of the given gap(s), the last n
492    outside those ranges.  Every offset of the network is eligible.
493
494    Raises ValueError if the constraints cannot be satisfied.
495    """
496    _check_version("random_ipv6s_with_range", network, 6)
497    return _random_addresses_with_range(network, gap, n, exclude_ips, exclude_nets)
498
499
500def _assign_from_topology(machine_nets, nets, ips, picker, nets_label):
501    """Fill the container *ips* from the networks of *nets* for every machine of *machine_nets*."""
502    assigned = {}  # {net_name: list[interface]} — used addresses per network, to avoid duplicates
503    for machine, net_names in machine_nets.items():
504        for net_name in net_names:
505            try:
506                network = getattr(nets, net_name)
507            except AttributeError:
508                raise AttributeError(
509                    f"random_ips_from_topology: data.{nets_label} has no network {net_name!r}") from None
510            ip = picker(network, 1, exclude_ips=assigned.get(net_name))[0]
511            assigned.setdefault(net_name, []).append(ip)
512            attr = machine if len(net_names) == 1 else f"{machine}_{net_name}"
513            setattr(ips, attr, ip)
514
515
516def random_ips_from_topology(data, topology, ipv4=True, ipv6=False):
517    """Assign random IPs to data.ips (from data.nets) and/or data.ips6 (from data.nets6)
518    based on a NetScheme0 topology.
519
520    topology: {net_name: [machine, ...] or {machine: iface_spec, ...}}
521    (the _topology class attribute format of NetScheme0)
522
523    For each machine m:
524    - belongs to exactly one network netX → data.ips.m  (in data.nets.netX)
525    - belongs to multiple networks       → data.ips.m_netX for each netX
526
527    `ipv4=True` (default) fills data.ips from data.nets as always; `ipv6=True` also fills
528    data.ips6 from data.nets6 with the same naming rule (use `ipv4=False, ipv6=True` for an
529    IPv6-only lab).  Every network of *topology* must exist in the containers used;
530    AttributeError otherwise.
531    """
532    if not ipv4 and not ipv6:
533        raise ValueError("random_ips_from_topology: at least one of ipv4/ipv6 must be True")
534    # Build inverse map: machine -> [net_name, ...]  (preserve insertion order)
535    machine_nets = {}
536    for net_name, machines in topology.items():
537        names = machines.keys() if isinstance(machines, dict) else machines
538        for m in names:
539            machine_nets.setdefault(m, []).append(net_name)
540
541    if ipv4:
542        _assign_from_topology(machine_nets, data.nets, data.ips, random_ipv4s, 'nets')
543    if ipv6:
544        _assign_from_topology(machine_nets, data.nets6, data.ips6, random_ipv6s, 'nets6')
545
546
547def random_mac_address(prefix=None, n=1):
548    """Return a list of n distinct random EUI MAC addresses.
549
550    `prefix` is an optional colon- or dash-separated hex string specifying the
551    leading bytes (e.g. ``"00:1A:2B"`` for a 3-byte OUI prefix).
552    Remaining bytes are chosen at random.
553
554    Raises ValueError if n distinct addresses cannot be generated.
555    """
556    if prefix is not None:
557        sep = "-" if "-" in prefix else ":"
558        prefix_bytes = bytes(int(x, 16) for x in prefix.split(sep))
559    else:
560        prefix_bytes = b""
561
562    suffix_len = 6 - len(prefix_bytes)
563    if suffix_len < 0:
564        raise ValueError(
565            f"Prefix too long: {prefix!r} ({len(prefix_bytes)} bytes, max 6)"
566        )
567
568    total = 256**suffix_len
569    if n > total:
570        raise ValueError(
571            f"Cannot generate {n} distinct MACs with {suffix_len} random bytes (max {total})"
572        )
573
574    def make(suffix_bytes):
575        all_bytes = prefix_bytes + suffix_bytes
576        if prefix is None:
577            all_bytes = bytes([all_bytes[0] & 0xFE]) + all_bytes[1:]
578        return EUI(":".join(f"{b:02x}" for b in all_bytes))
579
580    if total <= 65536:
581        candidates = [make(i.to_bytes(suffix_len, "big")) for i in range(total)]
582        return random.sample(candidates, n)
583
584    seen = set()
585    result = []
586    for _ in range(n):
587        for _ in range(1000):
588            suffix = random.randbytes(suffix_len)
589            if suffix not in seen:
590                seen.add(suffix)
591                result.append(make(suffix))
592                break
593        else:
594            raise ValueError(
595                f"Could not generate {n} distinct MACs after 1000 attempts"
596            )
597    return result