1 files changed, 33 insertions, 41 deletions

diff --git a/src/conf_mode/tftp_server.py b/src/conf_mode/tftp_server.py
index 32665ebc7..1a663a8be 100755
--- a/src/conf_mode/tftp_server.py
+++ b/src/conf_mode/tftp_server.py
@@ -23,64 +23,52 @@ from glob import glob
 from sys import exit
 
 from vyos.config import Config
-from vyos.validate import is_ipv4, is_addr_assigned
-from vyos import ConfigError
-from vyos.util import call
+from vyos.configdict import dict_merge
 from vyos.template import render
-
+from vyos.util import call
+from vyos.util import chmod_755
+from vyos.validate import is_ipv4
+from vyos.validate import is_addr_assigned
+from vyos.xml import defaults
+from vyos import ConfigError
 from vyos import airbag
 airbag.enable()
 
 config_file = r'/etc/default/tftpd'
 
-default_config_data = {
-    'directory': '',
-    'allow_upload': False,
-    'port': '69',
-    'listen': []
-}
-
 def get_config(config=None):
-    tftpd = deepcopy(default_config_data)
     if config:
         conf = config
     else:
         conf = Config()
+
     base = ['service', 'tftp-server']
     if not conf.exists(base):
         return None
-    else:
-        conf.set_level(base)
-
-    if conf.exists(['directory']):
-        tftpd['directory'] = conf.return_value(['directory'])
-
-    if conf.exists(['allow-upload']):
-        tftpd['allow_upload'] = True
-
-    if conf.exists(['port']):
-        tftpd['port'] = conf.return_value(['port'])
-
-    if conf.exists(['listen-address']):
-        tftpd['listen'] = conf.return_values(['listen-address'])
 
+    tftpd = conf.get_config_dict(base, key_mangling=('-', '_'), get_first_key=True)
+    # We have gathered the dict representation of the CLI, but there are default
+    # options which we need to update into the dictionary retrived.
+    default_values = defaults(base)
+    tftpd = dict_merge(default_values, tftpd)
     return tftpd
 
 def verify(tftpd):
     # bail out early - looks like removal from running config
-    if tftpd is None:
+    if not tftpd:
         return None
 
     # Configuring allowed clients without a server makes no sense
-    if not tftpd['directory']:
+    if 'directory' not in tftpd:
         raise ConfigError('TFTP root directory must be configured!')
 
-    if not tftpd['listen']:
+    if 'listen_address' not in tftpd:
         raise ConfigError('TFTP server listen address must be configured!')
 
-    for addr in tftpd['listen']:
-        if not is_addr_assigned(addr):
-            print('WARNING: TFTP server listen address {0} not assigned to any interface!'.format(addr))
+    for address in tftpd['listen_address']:
+        if not is_addr_assigned(address):
+            print(f'WARNING: TFTP server listen address "{address}" not ' \
+                  'assigned to any interface!')
 
     return None
 
@@ -95,15 +83,19 @@ def generate(tftpd):
         return None
 
     idx = 0
-    for listen in tftpd['listen']:
+    for address in tftpd['listen_address']:
         config = deepcopy(tftpd)
-        if is_ipv4(listen):
-            config['listen'] = [listen + ":" + tftpd['port'] + " -4"]
+        port = tftpd['port']
+        if is_ipv4(address):
+            config['listen_address'] = f'{address}:{port} -4'
         else:
-            config['listen'] = ["[" + listen + "]:" + tftpd['port'] + " -6"]
+            config['listen_address'] = f'[{address}]:{port} -6'
 
         file = config_file + str(idx)
-        render(file, 'tftp-server/default.tmpl', config)
+
+        import pprint
+        pprint.pprint(config)
+        render(file, 'tftp-server/default.tmpl', config, trim_blocks=True)
 
         idx = idx + 1
 
@@ -111,7 +103,7 @@ def generate(tftpd):
 
 def apply(tftpd):
     # stop all services first - then we will decide
-    call('systemctl stop tftpd@{0..20}.service')
+    call('systemctl stop tftpd@*.service')
 
     # bail out early - e.g. service deletion
     if tftpd is None:
@@ -120,7 +112,7 @@ def apply(tftpd):
     tftp_root = tftpd['directory']
     if not os.path.exists(tftp_root):
         os.makedirs(tftp_root)
-        os.chmod(tftp_root, stat.S_IRUSR|stat.S_IWUSR|stat.S_IXUSR|stat.S_IRGRP|stat.S_IXGRP|stat.S_IROTH|stat.S_IXOTH)
+        chmod_755(tftp_root)
 
     # get UNIX uid for user 'tftp'
     tftp_uid = pwd.getpwnam('tftp').pw_uid
@@ -135,8 +127,8 @@ def apply(tftpd):
         os.chown(tftp_root, tftp_uid, tftp_gid)
 
     idx = 0
-    for listen in tftpd['listen']:
-        call('systemctl restart tftpd@{0}.service'.format(idx))
+    for address in tftpd['listen_address']:
+        call(f'systemctl restart tftpd@{idx}.service')
         idx = idx + 1
 
     return None