#!/usr/bin/env python3
#
# Copyright (C) 2018 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
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# 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 jinja2
import vyos.validate

from vyos.config import Config
from vyos import ConfigError

config_file = r'/etc/default/tftpd-hpa'

# Please be careful if you edit the template.
config_tmpl = """
### Autogenerated by tftp_server.py ###

# See manual at https://linux.die.net/man/8/tftpd

TFTP_USERNAME="tftp"
TFTP_DIRECTORY="{{ directory }}"
{% if listen_ipv4 and listen_ipv6 -%}
TFTP_ADDRESS="{% for a in listen_ipv4 -%}{{ a }}:{{ port }}{{- " --address " if not loop.last -}}{% endfor -%} {% for a in listen_ipv6 %} --address [{{ a }}]:{{ port }}{% endfor -%}"
{% elif listen_ipv4 -%}
TFTP_ADDRESS="{% for a in listen_ipv4 -%}{{ a }}:{{ port }}{{- " --address " if not loop.last -}}{% endfor %} -4"
{% elif listen_ipv6 -%}
TFTP_ADDRESS="{% for a in listen_ipv6 -%}[{{ a }}]:{{ port }}{{- " --address " if not loop.last -}}{% endfor %} -6"
{%- endif %}

TFTP_OPTIONS="--secure {% if allow_upload %}--create --umask 000{% endif %}"

"""

default_config_data = {
    'directory': '',
    'allow_upload': False,
    'port': '69',
    'listen_ipv4': [],
    'listen_ipv6': []
}

def get_config():
    tftpd = default_config_data
    conf = Config()
    if not conf.exists('service tftp-server'):
        return None
    else:
        conf.set_level('service tftp-server')

    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'):
        for addr in conf.return_values('listen-address'):
            if vyos.validate.is_ipv4(addr):
                tftpd['listen_ipv4'].append(addr)
            else:
                tftpd['listen_ipv6'].append(addr)

    return tftpd

def verify(tftpd):
    # bail out early - looks like removal from running config
    if tftpd is None:
        return None

    # Configuring allowed clients without a server makes no sense
    if not tftpd['directory']:
        raise ConfigError('TFTP root directory must be configured!')

    if not (tftpd['listen_ipv4'] or tftpd['listen_ipv6']):
        raise ConfigError('TFTP server listen address must be configured!')

    for addr in tftpd['listen_ipv4']:
        # we always bind to localhost
        if '127.0.0.1' not in tftpd['listen_ipv4']:
            tftpd['listen_ipv4'].append('127.0.0.1')

        if not vyos.validate.is_addr_assigned(addr):
            print('WARNING: TFTP server listen address {0} not configured!'.format(addr))

    for addr in tftpd['listen_ipv6']:
        # we always bind to localhost
        if '::1' not in tftpd['listen_ipv6']:
            tftpd['listen_ipv6'].append('::1')

        if not vyos.validate.is_addr_assigned(addr):
            print('WARNING: TFTP server listen address {0} not configured!'.format(addr))

    return None

def generate(tftpd):
    # bail out early - looks like removal from running config
    if tftpd is None:
        return None

    tmpl = jinja2.Template(config_tmpl)
    config_text = tmpl.render(tftpd)
    with open(config_file, 'w') as f:
        f.write(config_text)

    return None

def apply(tftpd):
    if tftpd is not None:

        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)
            # get UNIX uid for user 'tftp'
            tftp_uid = pwd.getpwnam('tftp').pw_uid
            os.chown(tftp_root, tftp_uid, -1)

        os.system('sudo systemctl restart tftpd-hpa.service')
    else:
        # TFTP server support is removed in the commit
        os.system('sudo systemctl stop tftpd-hpa.service')
        os.unlink(config_file)

    return None

if __name__ == '__main__':
    try:
        c = get_config()
        verify(c)
        generate(c)
        apply(c)
    except ConfigError as e:
        print(e)
        sys.exit(1)