From b1cc7f721b686b09061bb484da2be03eb99da81e Mon Sep 17 00:00:00 2001
From: Ivan Pepelnjak <ip@ipspace.net>
Date: Sun, 11 Feb 2024 13:54:19 +0100
Subject: [PATCH] Snapshot #1

---
 netsim/devices/iosv.yml  |   6 ++
 netsim/devices/linux.yml |   5 ++
 netsim/modules/dhcp.py   | 153 +++++++++++++++++++++++++++++++++++++++
 netsim/modules/dhcp.yml  |  23 ++++++
 netsim/utils/strings.py  |  10 +++
 5 files changed, 197 insertions(+)
 create mode 100644 netsim/modules/dhcp.py
 create mode 100644 netsim/modules/dhcp.yml

diff --git a/netsim/devices/iosv.yml b/netsim/devices/iosv.yml
index 2f2057f18..5c5ab3803 100644
--- a/netsim/devices/iosv.yml
+++ b/netsim/devices/iosv.yml
@@ -30,6 +30,12 @@ features:
     community:
       standard: [ standard ]
       extended: [ extended ]
+  dhcp:
+    client:
+      ipv4: true
+      ipv6: true
+    relay: true
+    server: true
   initial:
     ipv4:
       unnumbered: false
diff --git a/netsim/devices/linux.yml b/netsim/devices/linux.yml
index 7aaaa4dbc..42976710b 100644
--- a/netsim/devices/linux.yml
+++ b/netsim/devices/linux.yml
@@ -2,6 +2,11 @@ description: Generic Linux host
 interface_name: eth{ifindex}
 mgmt_if: eth0
 role: host
+features:
+  dhcp:
+    client:
+      ipv4: true
+      ipv6: true
 libvirt:
   image: generic/ubuntu2004
   group_vars:
diff --git a/netsim/modules/dhcp.py b/netsim/modules/dhcp.py
new file mode 100644
index 000000000..f05033ba0
--- /dev/null
+++ b/netsim/modules/dhcp.py
@@ -0,0 +1,153 @@
+#
+# First-hop gateway transformation module
+#
+import typing
+from box import Box
+import netaddr
+
+from . import _Module,get_effective_module_attribute
+from ..utils import log, strings
+from .. import data
+from ..augment.nodes import reserve_id
+from ..augment import devices
+from ..data.validate import validate_attributes,must_be_string
+
+def check_protocol_support(node: Box, topology: Box) -> bool:
+  features = devices.get_device_features(node,topology.defaults)
+  OK = True
+
+  if node.get('dhcp.server',False) and not features.dhcp.server:
+    log.error(
+      f'Node {node.name} (device {node.device}) cannot be a DHCP server',
+      category=log.IncorrectValue,
+      module='dhcp')
+    OK = False
+
+  for intf in node.interfaces:
+    if not intf.get('dhcp.client',False):
+      continue
+    for af in ('ipv4','ipv6'):
+      if not af in intf.dhcp.client:
+        continue
+      if not features.dhcp.client[af]:
+        log.error(
+          f'Node {node.name} (device {node.device}) does not support {af} DHCP client',
+          more_data= [ 'DHCP client is used on interface {intf.ifname} ({intf.name})' ],
+          category=log.IncorrectValue,
+          module='dhcp')
+        OK = False
+  
+    for intf in node.interfaces:
+      if not intf.get('dhcp.server',False):
+        continue
+      if not features.dhcp.relay:
+        log.error(
+          f'Node {node.name} (device {node.device}) cannot be a DHCP relay',
+          category=log.IncorrectValue,
+          module='dhcp')
+        OK = False
+
+      if not topology.nodes.get(f'{intf.dhcp.server}.dhcp.server',False):
+        log.error(
+          f'Node {intf.dhcp.server} used for DHCP relaying on node {node.name} is not a DHCP server',
+          category=log.IncorrectValue,
+          module='dhcp')
+        OK = False
+        
+      if not intf.get('dhcp.vrf',False):
+        continue
+      if not features.dhcp.vrf:
+        log.error(
+          f'Node {node.name} (device {node.device}) cannot perform inter-VRF DHCP relaying',
+          category=log.IncorrectValue,
+          module='dhcp')
+        OK = False
+  
+      vrf = intf.get('dhcp.vrf')
+      if vrf == 'global':
+        continue
+
+      if vrf not in node.get('vrfs',{}):
+        log.error(
+          f'VRF {vrf} used for DHCP relaying is not used on node {node.name}',
+          category=log.IncorrectValue,
+          module='dhcp')
+        OK = False
+
+  return OK
+
+def build_topology_dhcp_pools(topology: Box) -> None:
+  topology.dhcp.pools = []
+
+  for link in topology.get('links',[]):
+    if not link.get('dhcp.subnet'):
+      continue
+
+    subnet = data.get_empty_box()
+
+    for af in ('ipv4','ipv6'):
+      if af not in link.dhcp.subnet or af not in link.prefix:
+        continue
+
+      subnet[af] = link.prefix[af]
+      subnet.name = link.get('name','') or link.get('_linkname','')
+      subnet.clean_name = strings.make_id(subnet.name)
+
+      if af in link.get('gateway'):
+        subnet.gateway[af] = link.gateway[af]
+
+      for intf in link.get('interfaces',[]):
+        if af not in intf or af in intf.get('dhcp.client',{}):
+          continue
+
+        if af not in subnet.excluded:
+          subnet.excluded[af] = []
+
+        addr = str(netaddr.IPNetwork(intf[af]).ip)
+        subnet.excluded[af].append(addr)
+        if af not in subnet.gateway:
+          subnet.gateway[af] = addr
+
+    topology.dhcp.pools.append(subnet)
+
+def set_dhcp_server_pools(node: Box, topology: Box) -> None:
+  if not node.get('dhcp.server',False):
+    return
+  
+  if not topology.get('dhcp.pools',False):
+    build_topology_dhcp_pools(topology)
+
+  node.dhcp.pools = topology.dhcp.pools
+
+class DHCP(_Module):
+
+  def link_pre_transform(self, link: Box, topology: Box) -> None:
+    for intf in link.get('interfaces',[]):
+      for af in ('ipv4','ipv6'):
+        if intf.get(af,False) != 'dhcp':
+          continue
+
+        if not intf.node in topology.nodes:
+          continue
+
+        intf.pop(af)
+        intf.dhcp.client[af] = True
+        link.dhcp.subnet[af] = True
+
+        node = topology.nodes[intf.node]
+        if 'module' not in node:
+          node.module = [ 'dhcp' ]
+        elif 'dhcp' not in node.module:
+          node.module.append('dhcp')
+
+  def node_post_transform(self, node: Box, topology: Box) -> None:
+    if not check_protocol_support(node,topology):
+      return
+
+    for intf in node.get('interfaces',[]):
+      for af in ('ipv4','ipv6'):
+        if intf.get(f'dhcp.client.{af}',False):
+          intf.pop(af,None)
+
+    if node.get('dhcp.server',False):
+      set_dhcp_server_pools(node,topology)
diff --git a/netsim/modules/dhcp.yml b/netsim/modules/dhcp.yml
new file mode 100644
index 000000000..fc4b74753
--- /dev/null
+++ b/netsim/modules/dhcp.yml
@@ -0,0 +1,23 @@
+# DHCP default settings and attributes
+#
+transform_after: [ vlan, vrf ]
+config_after: [ vlan, vrf ]
+attributes:
+  node:
+    server: bool
+  interface:
+    server: node_id
+    vrf: str
+    client:
+      ipv4: bool
+      ipv6: bool
+  link:
+    subnet:
+      ipv4: bool
+      ipv6: bool
+features:
+  ipv4: IPv4 DHCP client
+  ipv6: IPv6 DHCP client
+  relay: DHCP relay (IPv4 and IPv6)
+  server: DHCP server
+  vrf: Inter-VRF DHCP relay
diff --git a/netsim/utils/strings.py b/netsim/utils/strings.py
index 9f41cf424..e231b79f2 100644
--- a/netsim/utils/strings.py
+++ b/netsim/utils/strings.py
@@ -4,6 +4,8 @@
 import textwrap
 import typing
 import sys
+import re
+
 from box import Box,BoxList
 import rich.console, rich.table, rich.json, rich.syntax
 
@@ -167,3 +169,11 @@ def print_colored_text(txt: str, color: str, alt_txt: typing.Optional[str] = '',
     if alt_txt is not None:
       alt_txt = alt_txt or txt
       print(alt_txt,end='',file=sys.stderr if stderr else sys.stdout)
+
+"""
+make_id: Make an identifier out of a string
+"""
+def make_id(txt: str) -> str:
+  not_allowed = f'[^a-zA-Z0-9_]'
+  id = re.sub(not_allowed,'_',txt)
+  return id