#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#***************************************************************************
#                                  _   _ ____  _
#  Project                     ___| | | |  _ \| |
#                             / __| | | | |_) | |
#                            | (__| |_| |  _ <| |___
#                             \___|\___/|_| \_\_____|
#
# Copyright (C) Daniel Stenberg, <daniel@haxx.se>, et al.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at https://curl.se/docs/copyright.html.
#
# You may opt to use, copy, modify, merge, publish, distribute and/or sell
# copies of the Software, and permit persons to whom the Software is
# furnished to do so, under the terms of the COPYING file.
#
# This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY
# KIND, either express or implied.
#
# SPDX-License-Identifier: curl
#
###########################################################################
#
import difflib
import filecmp
import logging
import os
import shutil

import pytest
from testenv import CurlClient, Env, LocalClient, VsFTPD

log = logging.getLogger(__name__)


@pytest.mark.skipif(condition=not Env.has_vsftpd(), reason="missing vsftpd")
@pytest.mark.skipif(condition=Env.curl_uses_lib('rustls-ffi'),
                    reason="rustls does not support TLS session reuse")
@pytest.mark.skipif(condition=Env.curl_uses_lib('libressl'),
                    reason="libressl fails on TLS session reuse")
class TestVsFTPD:

    SUPPORTS_SSL = True

    @pytest.fixture(autouse=True, scope='class')
    def vsftpds(self, env):
        if not TestVsFTPD.SUPPORTS_SSL:
            pytest.skip('vsftpd does not seem to support SSL')
        vsftpds = VsFTPD(env=env, with_ssl=True)
        if not vsftpds.initial_start():
            vsftpds.stop()
            TestVsFTPD.SUPPORTS_SSL = False
            pytest.skip('vsftpd does not seem to support SSL')
        yield vsftpds
        vsftpds.stop()

    def _make_docs_file(self, docs_dir: str, fname: str, fsize: int):
        fpath = os.path.join(docs_dir, fname)
        data1k = 1024*'x'
        flen = 0
        with open(fpath, 'w') as fd:
            while flen < fsize:
                fd.write(data1k)
                flen += len(data1k)
        return flen

    @pytest.fixture(autouse=True, scope='class')
    def _class_scope(self, env, vsftpds):
        if os.path.exists(vsftpds.docs_dir):
            shutil.rmtree(vsftpds.docs_dir)
        if not os.path.exists(vsftpds.docs_dir):
            os.makedirs(vsftpds.docs_dir)
        self._make_docs_file(docs_dir=vsftpds.docs_dir, fname='data-1k', fsize=1024)
        self._make_docs_file(docs_dir=vsftpds.docs_dir, fname='data-10k', fsize=10 * 1024)
        self._make_docs_file(docs_dir=vsftpds.docs_dir, fname='data-1m', fsize=1024 * 1024)
        self._make_docs_file(docs_dir=vsftpds.docs_dir, fname='data-10m', fsize=10 * 1024 * 1024)
        env.make_data_file(indir=env.gen_dir, fname="upload-1k", fsize=1024)
        env.make_data_file(indir=env.gen_dir, fname="upload-100k", fsize=100 * 1024)
        env.make_data_file(indir=env.gen_dir, fname="upload-1m", fsize=1024 * 1024)

    def test_31_01_list_dir(self, env: Env, vsftpds: VsFTPD):
        curl = CurlClient(env=env)
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/'
        r = curl.ftp_ssl_get(urls=[url], with_stats=True)
        r.check_stats(count=1, http_status=226)
        with open(os.path.join(curl.run_dir, 'download_#1.data')) as fd:
            lines = fd.readlines()
        assert len(lines) == 4, f'list: {lines}'
        r.check_stats_timelines()

    # download 1 file, no SSL
    @pytest.mark.parametrize("docname", [
        'data-1k', 'data-1m', 'data-10m'
    ])
    def test_31_02_download_1(self, env: Env, vsftpds: VsFTPD, docname):
        curl = CurlClient(env=env)
        srcfile = os.path.join(vsftpds.docs_dir, f'{docname}')
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}?[0-{count-1}]'
        r = curl.ftp_ssl_get(urls=[url], with_stats=True)
        r.check_stats(count=count, http_status=226)
        self.check_downloads(curl, srcfile, count)
        r.check_stats_timelines()

    @pytest.mark.parametrize("docname,count,secure", [
        ['data-1k', 10, True],
        ['data-1m', 5, True],
        ['data-1m', 5, False],
        ['data-10m', 2,True]
    ])
    def test_31_03_download_10_serial(self, env: Env, vsftpds: VsFTPD, docname, count, secure):
        curl = CurlClient(env=env)
        srcfile = os.path.join(vsftpds.docs_dir, f'{docname}')
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}?[0-{count-1}]'
        xargs = []
        if not secure:
            xargs.append('--insecure')
        r = curl.ftp_ssl_get(urls=[url], with_stats=True, extra_args=xargs)
        r.check_stats(count=count, http_status=226)
        self.check_downloads(curl, srcfile, count)
        assert r.total_connects == count + 1, 'should reuse the control conn'
        r.check_stats_timelines()

    @pytest.mark.parametrize("docname,count,secure", [
        ['data-1k', 10, True],
        ['data-1m', 5, True],
        ['data-1m', 5, False],
        ['data-10m', 2,True]
    ])
    def test_31_04_download_10_parallel(self, env: Env, vsftpds: VsFTPD, docname, count, secure):
        curl = CurlClient(env=env)
        srcfile = os.path.join(vsftpds.docs_dir, f'{docname}')
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}?[0-{count-1}]'
        xargs = ['--parallel']
        if not secure:
            xargs.append('--insecure')
        r = curl.ftp_ssl_get(urls=[url], with_stats=True, extra_args=xargs)
        r.check_stats(count=count, http_status=226)
        self.check_downloads(curl, srcfile, count)
        assert r.total_connects > count + 1, 'should have used several control conns'
        r.check_stats_timelines()

    @pytest.mark.parametrize("docname", [
        'upload-1k', 'upload-100k', 'upload-1m'
    ])
    def test_31_05_upload_1(self, env: Env, vsftpds: VsFTPD, docname):
        curl = CurlClient(env=env)
        srcfile = os.path.join(env.gen_dir, docname)
        dstfile = os.path.join(vsftpds.docs_dir, docname)
        self._rmf(dstfile)
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/'
        r = curl.ftp_ssl_upload(urls=[url], fupload=f'{srcfile}', with_stats=True)
        r.check_stats(count=count, http_status=226)
        self.check_upload(env, vsftpds, docname=docname)
        r.check_stats_timelines()

    def _rmf(self, path):
        if os.path.exists(path):
            os.remove(path)

    # check with `tcpdump` if curl causes any TCP RST packets
    @pytest.mark.skipif(condition=not Env.tcpdump(), reason="tcpdump not available")
    @pytest.mark.skipif(condition=not Env.curl_is_debug(), reason="needs curl debug")
    @pytest.mark.skipif(condition=not Env.curl_is_verbose(), reason="needs curl verbose strings")
    def test_31_06_shutdown_download(self, env: Env, vsftpds: VsFTPD):
        docname = 'data-1k'
        curl = CurlClient(env=env)
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}?[0-{count-1}]'
        r = curl.ftp_ssl_get(urls=[url], with_stats=True, with_tcpdump=True)
        r.check_stats(count=count, http_status=226)
        # vsftp closes control connection without niceties,
        # look only at ports from DATA connection.
        data_ports = vsftpds.get_data_ports(r)
        assert len(data_ports), f'unable to find FTP data port connected to\n{r.dump_logs()}'
        assert len(r.tcpdump.get_rsts(port_pairs=data_ports)) == 0, 'Unexpected TCP RST packets'

    # check with `tcpdump` if curl causes any TCP RST packets
    @pytest.mark.skipif(condition=not Env.tcpdump(), reason="tcpdump not available")
    @pytest.mark.skipif(condition=not Env.curl_is_debug(), reason="needs curl debug")
    @pytest.mark.skipif(condition=not Env.curl_is_verbose(), reason="needs curl verbose strings")
    def test_31_07_shutdown_upload(self, env: Env, vsftpds: VsFTPD):
        docname = 'upload-1k'
        curl = CurlClient(env=env)
        srcfile = os.path.join(env.gen_dir, docname)
        dstfile = os.path.join(vsftpds.docs_dir, docname)
        self._rmf(dstfile)
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/'
        r = curl.ftp_ssl_upload(urls=[url], fupload=f'{srcfile}', with_stats=True, with_tcpdump=True)
        r.check_stats(count=count, http_status=226)
        # vsftp closes control connection without niceties,
        # look only at ports from DATA connection.
        data_ports = vsftpds.get_data_ports(r)
        assert len(data_ports), f'unable to find FTP data port connected to\n{r.dump_logs()}'
        assert len(r.tcpdump.get_rsts(port_pairs=data_ports)) == 0, 'Unexpected TCP RST packets'

    def test_31_08_upload_ascii(self, env: Env, vsftpds: VsFTPD):
        docname = 'upload-ascii'
        line_length = 21
        srcfile = os.path.join(env.gen_dir, docname)
        dstfile = os.path.join(vsftpds.docs_dir, docname)
        env.make_data_file(indir=env.gen_dir, fname=docname, fsize=100 * 1024,
                           line_length=line_length)
        srcsize = os.path.getsize(srcfile)
        self._rmf(dstfile)
        count = 1
        curl = CurlClient(env=env)
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/'
        r = curl.ftp_ssl_upload(urls=[url], fupload=f'{srcfile}', with_stats=True,
                                extra_args=['--use-ascii'])
        r.check_stats(count=count, http_status=226)
        # expect the uploaded file to be number of converted newlines larger
        dstsize = os.path.getsize(dstfile)
        with open(srcfile) as fd:
            newlines = len(fd.readlines())
        assert (srcsize + newlines) == dstsize, \
            f'expected source with {newlines} lines to be that much larger,'\
            f'instead srcsize={srcsize}, upload size={dstsize}, diff={dstsize-srcsize}'

    def test_31_08_active_download(self, env: Env, vsftpds: VsFTPD):
        docname = 'data-10k'
        curl = CurlClient(env=env)
        srcfile = os.path.join(vsftpds.docs_dir, f'{docname}')
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}?[0-{count-1}]'
        r = curl.ftp_ssl_get(urls=[url], with_stats=True, extra_args=[
            '--ftp-port', '127.0.0.1'
        ])
        r.check_stats(count=count, http_status=226)
        self.check_downloads(curl, srcfile, count)

    def test_31_09_active_upload(self, env: Env, vsftpds: VsFTPD):
        docname = 'upload-1k'
        curl = CurlClient(env=env)
        srcfile = os.path.join(env.gen_dir, docname)
        dstfile = os.path.join(vsftpds.docs_dir, docname)
        self._rmf(dstfile)
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/'
        r = curl.ftp_ssl_upload(urls=[url], fupload=f'{srcfile}', with_stats=True, extra_args=[
            '--ftp-port', '127.0.0.1'
        ])
        r.check_stats(count=count, http_status=226)
        self.check_upload(env, vsftpds, docname=docname)

    @pytest.mark.parametrize("indata", [
        pytest.param('1234567890', id='10-bytes'),
        pytest.param('', id='0-bytes'),
    ])
    def test_31_10_upload_stdin(self, env: Env, vsftpds: VsFTPD, indata):
        curl = CurlClient(env=env)
        docname = "upload_31_10"
        dstfile = os.path.join(vsftpds.docs_dir, docname)
        self._rmf(dstfile)
        count = 1
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}'
        r = curl.ftp_ssl_upload(urls=[url], updata=indata, with_stats=True)
        r.check_stats(count=count, http_status=226)
        assert os.path.exists(dstfile)
        with open(dstfile) as fd:
            destdata = fd.readlines()
        expdata = [indata] if len(indata) else []
        assert expdata == destdata, f'expected: {expdata}, got: {destdata}'

    def test_31_11_download_non_existing(self, env: Env, vsftpds: VsFTPD):
        curl = CurlClient(env=env)
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/does-not-exist'
        r = curl.ftp_ssl_get(urls=[url], with_stats=True)
        r.check_exit_code(78)
        r.check_stats(count=1, exitcode=78)

    def test_31_12_upload_eprt(self, env: Env, vsftpds: VsFTPD):
        docname = 'test_31_12'
        client = LocalClient(name='cli_ftp_upload', env=env)
        if not client.exists():
            pytest.skip(f'example client not built: {client.name}')
        url = f'ftp://{env.ftp_domain}:{vsftpds.port}/{docname}'
        r = client.run(args=['-r', f'{env.ftp_domain}:{vsftpds.port}:127.0.0.1', url])
        r.check_exit_code(0)
        dstfile = os.path.join(vsftpds.docs_dir, docname)
        assert os.path.exists(dstfile), f'{r.dump_logs()}'

    # connection reuse with STARTTLS required
    # 1st download without STARTTLS, 2nd with --ssl-reqd
    @pytest.mark.skipif(condition=not Env.curl_is_debug(), reason="needs curl debug")
    def test_31_13_starttls_reuse(self, env: Env, vsftpds: VsFTPD):
        run_env = os.environ.copy()
        run_env['CURL_DBG_NO_USE_SSL_ON_FIRST'] = '1'
        curl = CurlClient(env=env, run_env=run_env)
        url1 = f'ftp://{env.ftp_domain}:{vsftpds.port}/data-1k'
        url2 = f'ftp://{env.ftp_domain}:{vsftpds.port}/data-10k'
        r = curl.run_direct(with_stats=True, args=[
            '-svv', '--resolve', f'{env.ftp_domain}:{vsftpds.port}:127.0.0.1',
            '--cacert', env.ca.cert_file,
            url1, '--out-null',
            url2, '--out-null', '--ssl-reqd'
        ])
        r.check_exit_code(0)
        r.check_stats(count=2, http_status=226)
        # expect 4 connections to have been made:
        # 1. 1st CONTROL without STARTTLS
        # 2. 1st DATA for download
        # 3. 2nd CONTROL with STARTTLS (not reuse of 1)
        # 4. 2nd DATA for download
        assert r.total_connects == 4, f'{r.dump_logs()}'

    def check_downloads(self, client, srcfile: str, count: int,
                        complete: bool = True):
        for i in range(count):
            dfile = client.download_file(i)
            assert os.path.exists(dfile)
            if complete and not filecmp.cmp(srcfile, dfile, shallow=False):
                with open(srcfile) as fa, open(dfile) as fb:
                    a = fa.readlines()
                    b = fb.readlines()
                diff = "".join(difflib.unified_diff(a=a, b=b,
                                                    fromfile=srcfile,
                                                    tofile=dfile,
                                                    n=1))
                assert False, f'download {dfile} differs:\n{diff}'

    def check_upload(self, env, vsftpd: VsFTPD, docname):
        srcfile = os.path.join(env.gen_dir, docname)
        dstfile = os.path.join(vsftpd.docs_dir, docname)
        assert os.path.exists(srcfile)
        assert os.path.exists(dstfile)
        if not filecmp.cmp(srcfile, dstfile, shallow=False):
            with open(srcfile) as fa, open(dstfile) as fb:
                a = fa.readlines()
                b = fb.readlines()
            diff = "".join(difflib.unified_diff(a=a, b=b,
                                                fromfile=srcfile,
                                                tofile=dstfile,
                                                n=1))
            assert False, f'upload {dstfile} differs:\n{diff}'
