1 files changed, 33 insertions, 33 deletions

diff --git a/src/conf_mode/tftp_server.py b/src/conf_mode/tftp_server.py
index ff7cad0c9..7a7246783 100755
--- a/src/conf_mode/tftp_server.py
+++ b/src/conf_mode/tftp_server.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 #
-# Copyright (C) 2018 VyOS maintainers and contributors
+# Copyright (C) 2018-2020 VyOS maintainers and contributors
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 2 or later as
@@ -13,31 +13,24 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
-#
-#
 
-import sys
 import os
 import stat
 import pwd
-import copy
-import glob
 
-import jinja2
-import vyos.validate
+from copy import deepcopy
+from glob import glob
+from jinja2 import FileSystemLoader, Environment
+from sys import exit
 
 from vyos.config import Config
+from vyos.defaults import directories as vyos_data_dir
+from vyos.validate import is_ipv4, is_addr_assigned
 from vyos import ConfigError
+from vyos.util import call
 
-config_file = r'/etc/default/tftpd'
-
-# Please be careful if you edit the template.
-config_tmpl = """
-### Autogenerated by tftp_server.py ###
-DAEMON_ARGS="--listen --user tftp --address {% for a in listen-%}{{ a }}{% endfor %}{% if allow_upload %} --create --umask 000{% endif %} --secure {{ directory }}"
 
-
-"""
+config_file = r'/etc/default/tftpd'
 
 default_config_data = {
     'directory': '',
@@ -47,23 +40,25 @@ default_config_data = {
 }
 
 def get_config():
-    tftpd = copy.deepcopy(default_config_data)
+    tftpd = deepcopy(default_config_data)
     conf = Config()
-    if not conf.exists('service tftp-server'):
+    base = ['service', 'tftp-server']
+    if not conf.exists(base):
         return None
     else:
-        conf.set_level('service tftp-server')
+        conf.set_level(base)
 
-    if conf.exists('directory'):
-        tftpd['directory'] = conf.return_value('directory')
+    if conf.exists(['directory']):
+        tftpd['directory'] = conf.return_value(['directory'])
 
-    if conf.exists('allow-upload'):
+    if conf.exists(['allow-upload']):
         tftpd['allow_upload'] = True
 
-    if conf.exists('port'):
-        tftpd['port'] = conf.return_value('port')
+    if conf.exists(['port']):
+        tftpd['port'] = conf.return_value(['port'])
 
-    tftpd['listen'] = conf.return_values('listen-address')
+    if conf.exists(['listen-address']):
+        tftpd['listen'] = conf.return_values(['listen-address'])
 
     return tftpd
 
@@ -80,7 +75,7 @@ def verify(tftpd):
         raise ConfigError('TFTP server listen address must be configured!')
 
     for addr in tftpd['listen']:
-        if not vyos.validate.is_addr_assigned(addr):
+        if not is_addr_assigned(addr):
             print('WARNING: TFTP server listen address {0} not assigned to any interface!'.format(addr))
 
     return None
@@ -88,22 +83,27 @@ def verify(tftpd):
 def generate(tftpd):
     # cleanup any available configuration file
     # files will be recreated on demand
-    for i in glob.glob(config_file + '*'):
+    for i in glob(config_file + '*'):
         os.unlink(i)
 
     # bail out early - looks like removal from running config
     if tftpd is None:
         return None
 
+    # Prepare Jinja2 template loader from files
+    tmpl_path = os.path.join(vyos_data_dir['data'], 'templates', 'tftp-server')
+    fs_loader = FileSystemLoader(tmpl_path)
+    env = Environment(loader=fs_loader)
+
     idx = 0
     for listen in tftpd['listen']:
-        config = copy.deepcopy(tftpd)
-        if vyos.validate.is_ipv4(listen):
+        config = deepcopy(tftpd)
+        if is_ipv4(listen):
             config['listen'] = [listen + ":" + tftpd['port'] + " -4"]
         else:
             config['listen'] = ["[" + listen + "]" + tftpd['port'] + " -6"]
 
-        tmpl = jinja2.Template(config_tmpl)
+        tmpl = env.get_template('default.tmpl')
         config_text = tmpl.render(config)
         file = config_file + str(idx)
         with open(file, 'w') as f:
@@ -115,7 +115,7 @@ def generate(tftpd):
 
 def apply(tftpd):
     # stop all services first - then we will decide
-    os.system('sudo systemctl stop tftpd@{0..20}')
+    call('systemctl stop tftpd@{0..20}')
 
     # bail out early - e.g. service deletion
     if tftpd is None:
@@ -140,7 +140,7 @@ def apply(tftpd):
 
     idx = 0
     for listen in tftpd['listen']:
-        os.system('sudo systemctl restart tftpd@{0}.service'.format(idx))
+        call('systemctl restart tftpd@{0}.service'.format(idx))
         idx = idx + 1
 
     return None
@@ -153,4 +153,4 @@ if __name__ == '__main__':
         apply(c)
     except ConfigError as e:
         print(e)
-        sys.exit(1)
+        exit(1)