mdns_repeater: cleanup python implementation

1 files changed, 55 insertions, 37 deletions

diff --git a/src/conf_mode/mdns_repeater.py b/src/conf_mode/mdns_repeater.py
index 474a6a5cf..cef735c0d 100755
--- a/src/conf_mode/mdns_repeater.py
+++ b/src/conf_mode/mdns_repeater.py
@@ -18,7 +18,7 @@
 
 import sys
 import os
-
+import jinja2
 import netifaces
 
 from vyos.config import Config
@@ -26,60 +26,78 @@ from vyos import ConfigError
 
 config_file = r'/etc/default/mdns-repeater'
 
-def get_config():
-    interface_list = []
+config_tmpl = """
+### Autogenerated by mdns_repeater.py ###
+DAEMON_ARGS="{{ interfaces | join(' ') }}"
+"""
+
+default_config_data = {
+    'disabled': False,
+    'interfaces': []
+}
 
+def get_config():
+    mdns = default_config_data
     conf = Config()
-    conf.set_level('service mdns repeater')
-    if not conf.exists(''):
-        return interface_list
+    if not conf.exists('service mdns repeater'):
+        return None
+    else:
+        conf.set_level('service mdns repeater')
 
-    if conf.exists('interface'):
-        intfs_names = []
-        intfs_names = conf.return_values('interface')
+    # Service can be disabled by user
+    if conf.exists('disable'):
+        mdns['disabled'] = True
+        return mdns
 
-        for name in intfs_names:
-            interface_list.append(name)
+    # Interface to repeat mDNS advertisements
+    if conf.exists('interface'):
+        mdns['interfaces'] = conf.return_values('interface')
 
-    return interface_list
+    return mdns
 
 def verify(mdns):
-    # '0' interfaces are possible, think of service deletion. Only '1' is not supported!
-    if len(mdns) == 1:
-        raise ConfigError('At least 2 interfaces must be specified but %d given!' % len(mdns))
-
-    # For mdns-repeater to work it is essential that the interfaces
-    # have an IP address assigned
-    for intf in mdns:
-        try:
-            netifaces.ifaddresses(intf)[netifaces.AF_INET]
-        except KeyError as e:
-            raise ConfigError('No IP address configured for interface "%s"!' % intf)
+    if mdns is None:
+        return None
+
+    if mdns['disabled']:
+        return None
+
+    # We need at least two interfaces to repeat mDNS advertisments
+    if len(mdns['interfaces']) < 2:
+        raise ConfigError('mDNS repeater requires at least 2 configured interfaces!')
+
+    # For mdns-repeater to work it is essential that the interfaces has
+    # an IPv4 address assigned
+    for interface in mdns['interfaces']:
+        if netifaces.AF_INET in netifaces.ifaddresses(interface).keys():
+            if len(netifaces.ifaddresses(interface)[netifaces.AF_INET]) < 1:
+                raise ConfigError('mDNS repeater requires an IPv6 address configured on interface %s!'.format(interface))
 
     return None
 
 def generate(mdns):
-    config_header = '### Autogenerated by mdns_repeater.py ###\n'
-    if len(mdns) > 0:
-        config_args = 'DAEMON_ARGS="' + ' '.join(str(e) for e in mdns) + '"\n'
-    else:
-        config_args = 'DAEMON_ARGS=""\n'
+    if mdns is None:
+        return None
+
+    if mdns['disabled']:
+        print('Warning: mDNS repeater will be deactivated because it is disabled')
+        return None
 
-    # write new configuration file
-    f = open(config_file, 'w')
-    f.write(config_header)
-    f.write(config_args)
-    f.close()
+    tmpl = jinja2.Template(config_tmpl)
+    config_text = tmpl.render(mdns)
+    with open(config_file, 'w') as f:
+        f.write(config_text)
 
     return None
 
 def apply(mdns):
-    if len(mdns) == 0:
-        cmd = "sudo systemctl stop mdns-repeater"
+    if (mdns is None) or mdns['disabled']:
+        os.system('sudo systemctl stop mdns-repeater')
+        if os.path.exists(config_file):
+            os.unlink(config_file)
     else:
-        cmd = "sudo systemctl restart mdns-repeater"
+        os.system('sudo systemctl restart mdns-repeater')
 
-    os.system(cmd)
     return None
 
 if __name__ == '__main__':