# Copyright (C) 2024–present  Loren Eteval & contributors <loren.eteval@proton.me>
#
# This file is part of Furious.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# 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 <https://www.gnu.org/licenses/>.

"""Verify sing-tun's Linux helpers without changing real host networking."""

from Furious.Models.SingTUN import prepareSingTUNSettings, addSingTUNExclusions
from Furious.Service.SingTUNHost import SingTUNHostPlan
from Furious.Frozenlib.SingTUNRouting import (
    SingTUNLinuxRouting,
    _linuxRouteRanges,
    _linuxRules,
    _linuxRuleDelete,
    _linuxRuleSelectors,
)

from pathlib import Path
from unittest import mock

import importlib
import ipaddress
import json
import os
import shutil
import shlex
import subprocess
import tempfile
import threading
import unittest


def _bash():
    """Use native Bash or Git Bash, never a Windows WSL launcher."""
    if os.name != 'nt':
        return shutil.which('bash')

    git = shutil.which('git')

    if git:
        candidate = Path(git).resolve().parent.parent / 'bin' / 'bash.exe'

        if candidate.exists():
            return str(candidate)

    return None


class SingTUNLinuxHostTest(unittest.TestCase):
    def setUp(self):
        self.module = importlib.import_module('Furious.Frozenlib.SingTUNRouting')
        self.tun = prepareSingTUNSettings(
            {'tun_options': {'Name': 'utun101', 'Inet6Address': ['fd00::1/64']}}
        )['tun_options']
        self.host = SingTUNLinuxRouting(self.tun, 123456, 2000, threading.Event())

    def tearDown(self):
        if self.host._process is not None:
            self.host._process.poll.return_value = 0

        self.host.drain()

    def testHelperCreatesUserOwnedTunAndQuotesOnlyHostArguments(self):
        scripts = []
        process = mock.Mock()
        process.wait.return_value = 0
        process.poll.return_value = 0

        def launch(arguments, **kwargs):
            self.assertEqual(arguments[:2], ['pkexec', 'bash'])
            self.assertEqual(kwargs['stdin'], subprocess.DEVNULL)
            self.assertTrue(
                Path(arguments[2]).name.startswith(tempfile.gettempprefix())
            )
            self.assertEqual(Path(arguments[2]).suffix, '.sh')
            scripts.append(Path(arguments[2]).read_text(encoding='utf-8'))
            return process

        link = {'ifname': 'utun101', 'ifindex': 77}

        with mock.patch.object(
            self.module.os, 'geteuid', return_value=1000, create=True
        ), mock.patch.object(
            self.module.os, 'getuid', return_value=1000, create=True
        ), mock.patch.object(
            self.module.subprocess, 'Popen', side_effect=launch
        ), mock.patch.object(
            self.module, '_command', return_value=json.dumps([link])
        ):
            self.host.prepare()
            self.host.applyDNS('1.1.1.1')

        self.assertEqual(self.host.index, 77)
        self.assertIn('ip tuntap add dev utun101 mode tun user 1000', scripts[0])
        self.assertIn('ip link set dev utun101 mtu 1500', scripts[0])
        self.assertIn('ip -6 addr add fd00::1/64 dev utun101', scripts[0])
        self.assertIn('ip link set dev utun101 up', scripts[0])
        self.assertIn('createdIndex=', scripts[0])
        self.assertIn('trap ', scripts[0])
        self.assertNotIn('ifalias', ''.join(scripts))
        self.assertNotIn('furious-sing-tun-', ''.join(scripts))
        self.assertIn('/sys/class/net/utun101/ifindex', scripts[1])
        self.assertIn('resolvectl dns utun101 1.1.1.1', scripts[1])
        self.assertNotIn('socks5', ''.join(scripts))
        self.assertIsNone(self.host._process)
        self.assertIsNone(self.host._directory)

        bash = _bash()

        if bash and Path(bash).exists():
            for script in scripts:
                checked = subprocess.run(
                    [bash, '-n', '-c', script], capture_output=True, timeout=5
                )
                self.assertEqual(
                    checked.returncode, 0, checked.stderr.decode(errors='replace')
                )

    def testAuthorizationTimeoutRetainsExactProcessAndCancellationFile(self):
        process = mock.Mock()
        process.wait.side_effect = subprocess.TimeoutExpired('pkexec', 60)
        process.poll.return_value = None

        with mock.patch.object(
            self.module.os, 'geteuid', return_value=1000, create=True
        ), mock.patch.object(self.module.subprocess, 'Popen', return_value=process):
            with self.assertRaisesRegex(RuntimeError, 'timed out'):
                self.host.execute(
                    [['ip', 'link', 'set', 'dev', 'utun101', 'up']], preparing=True
                )

        directory = Path(self.host._directory.name)
        self.assertIs(self.host._process, process)
        self.assertEqual(len(list(directory.glob('*.sh'))), 1)
        self.assertTrue((directory / 'cancelled').exists())

        with self.assertRaisesRegex(RuntimeError, 'still running'):
            self.host.drain()

        self.assertIs(self.host._process, process)
        process.terminate.assert_not_called()
        process.kill.assert_not_called()

        process.poll.return_value = 0
        self.host.drain()

        self.assertFalse(directory.exists())
        self.assertIsNone(self.host._process)

    def testScriptRetainsSuccessfulSetupAndRollsBackFailure(self):
        bash = _bash()

        if not bash or not Path(bash).exists():
            self.skipTest('Bash is unavailable for the isolated script fixture')

        realPopen = subprocess.Popen

        failures = (
            ('mtu', 'cancel', 'abandon', 'none')
            if os.name != 'nt'
            else ('mtu', 'cancel', 'none')
        )

        for failure in failures:
            with self.subTest(
                failure=failure
            ), tempfile.TemporaryDirectory() as directory:
                root = Path(directory).resolve()
                net = root / 'net'
                binDirectory = root / 'bin'
                net.mkdir()
                binDirectory.mkdir()
                ip = binDirectory / 'ip'

                # This fixture only creates/removes files under its temporary
                # root. No real ip, PolicyKit, TUN, route or DNS command is run.
                with ip.open('w', encoding='utf-8', newline='\n') as file:
                    file.write(
                        '#!/usr/bin/env bash\n'
                        'set -eu\n'
                        'device="$FAKE_NET/utun101"\n'
                        'case "$1 $2 $3" in\n'
                        '"tuntap add dev")\n'
                        '  mkdir "$device"\n'
                        '  echo 77 > "$device/ifindex"\n'
                        '  if [ "$FAKE_FAILURE" = cancel ]; then touch "$FAKE_CANCEL"; fi\n'
                        '  if [ "$FAKE_FAILURE" = abandon ]; then\n'
                        '    rm "$FAKE_SCRIPT"\n'
                        '    rmdir "${FAKE_CANCEL%/cancelled}"\n'
                        '  fi\n'
                        '  ;;\n'
                        '"link set dev")\n'
                        '  if [ "$FAKE_FAILURE" = mtu ] && [ "$5" = mtu ]; then exit 1; fi\n'
                        '  ;;\n'
                        '"-4 addr add"|"-6 addr add"|"-4 route add"|"-6 route add"|"-4 rule add"|"-6 rule add") ;;\n'
                        '"link delete dev")\n'
                        '  rm "$device/ifindex"\n'
                        '  rmdir "$device"\n'
                        '  ;;\n'
                        '*) exit 90 ;;\n'
                        'esac\n'
                    )

                ip.chmod(0o755)

                def launch(arguments, **kwargs):
                    self.assertEqual(arguments[:2], ['pkexec', 'bash'])
                    script = Path(arguments[-1])
                    text = script.read_text(encoding='utf-8')
                    text = text.replace('/sys/class/net', net.as_posix())
                    cancelled = script.parent / 'cancelled'
                    text = text.replace(str(cancelled), cancelled.as_posix())
                    text = text.replace(str(script.parent), script.parent.as_posix())

                    binPath = shlex.quote(binDirectory.as_posix())

                    if os.name == 'nt':
                        binPath = '$(cygpath -u ' + binPath + ')'

                    with script.open('w', encoding='utf-8', newline='\n') as file:
                        file.write('PATH=' + binPath + ':"$PATH"\n' + text)

                    environment = os.environ.copy()
                    environment.update(
                        FAKE_NET=net.as_posix(),
                        FAKE_FAILURE=failure,
                        FAKE_CANCEL=(script.parent / 'cancelled').as_posix(),
                        FAKE_SCRIPT=script.as_posix(),
                    )
                    return realPopen([bash, str(script)], env=environment, **kwargs)

                with mock.patch.object(
                    self.module.os, 'geteuid', return_value=1000, create=True
                ), mock.patch.object(
                    self.module.os, 'getuid', return_value=1000, create=True
                ), mock.patch.object(
                    self.module.subprocess, 'Popen', side_effect=launch
                ), mock.patch.object(
                    self.module,
                    '_command',
                    return_value=json.dumps([{'ifname': 'utun101', 'ifindex': 77}]),
                ):
                    if failure == 'none':
                        self.host.prepare()
                        self.assertEqual(self.host.index, 77)
                    else:
                        with self.assertRaisesRegex(
                            RuntimeError, 'host operation failed'
                        ):
                            self.host.prepare()

                self.assertEqual((net / 'utun101').exists(), failure == 'none')
                self.assertIsNone(self.host._process)
                self.assertIsNone(self.host._directory)

    def testCancellationBeforeAuthorizationCannotAcquireDevice(self):
        scripts = []
        process = mock.Mock()
        process.wait.return_value = 1
        process.poll.return_value = 1

        def launch(arguments, **kwargs):
            directory = Path(arguments[-1]).parent
            self.host.cancelled.set()
            self.host.cancel()
            self.assertTrue((directory / 'cancelled').exists())
            scripts.append(Path(arguments[-1]).read_text(encoding='utf-8'))
            return process

        with mock.patch.object(
            self.module.os, 'geteuid', return_value=1000, create=True
        ), mock.patch.object(self.module.subprocess, 'Popen', side_effect=launch):
            with self.assertRaisesRegex(RuntimeError, 'authorization was cancelled'):
                self.host.execute(
                    [['ip', 'link', 'set', 'dev', 'utun101', 'up']], preparing=True
                )

        self.assertLess(scripts[0].index('cancelled'), scripts[0].index('ip link'))
        self.assertIsNone(self.host._process)

    def testCancelledWorkDoesNotLaunchAnotherAuthorization(self):
        self.host.cancelled.set()

        with mock.patch.object(self.module.subprocess, 'Popen') as launch:
            with self.assertRaisesRegex(RuntimeError, 'preparation was cancelled'):
                self.host.execute(
                    [['ip', 'link', 'set', 'dev', 'utun101', 'up']], preparing=True
                )

        launch.assert_not_called()
        self.assertIsNone(self.host._directory)

    def testRootHelperRunsWithoutPolicyKitAndFailedLaunchReleasesScript(self):
        with mock.patch.object(
            self.module.os, 'geteuid', return_value=0, create=True
        ), mock.patch.object(
            self.module.subprocess, 'Popen', side_effect=OSError('launch')
        ) as launch:
            with self.assertRaisesRegex(RuntimeError, 'Unable to start'):
                self.host.execute([['ip', 'link', 'set', 'dev', 'utun101', 'up']])

        self.assertEqual(launch.call_args.args[0][0], 'bash')
        self.assertIsNone(self.host._directory)

    def testRouteRangesKeepCustomPrefixesAndAllRemoteExclusions(self):
        configuration = addSingTUNExclusions(
            {'tun_options': self.tun}, ['192.0.2.8', '2001:db8::8']
        )
        ranges = _linuxRouteRanges(configuration['tun_options'])

        for address in ('192.0.2.8', '2001:db8::8'):
            ip = ipaddress.ip_address(address)
            self.assertFalse(
                any(ip.version == prefix.version and ip in prefix for prefix in ranges)
            )

        self.tun['Inet4RouteAddress'] = ['203.0.113.0/24']
        ranges = _linuxRouteRanges(self.tun)
        self.assertEqual(
            [prefix for prefix in ranges if prefix.version == 4],
            [ipaddress.ip_network('203.0.113.0/24')],
        )

    def testIPv6LoopPreventionPrecedesSourceAndDefaultRules(self):
        rules = _linuxRules(self.tun, 123456, 2000)
        v6 = [command for command in rules if command[1] == '-6']
        goto = next(command for command in v6 if 'utun101' in command)
        source = next(command for command in v6 if 'fd00::/64' in command)
        default = next(command for command in v6 if command[6:] == ['table', '123456'])

        self.assertEqual(goto[5], '2001')
        self.assertEqual(goto[-2:], ['goto', '2010'])
        self.assertEqual(source[5], '2002')
        self.assertEqual(default[5], '2004')

    def testReplacementInterfaceCannotReceiveDnsOrBeDeleted(self):
        self.host.index = 77
        link = {'ifname': 'utun101', 'ifindex': 78}

        with mock.patch.object(
            self.module, '_command', return_value=json.dumps([link])
        ), mock.patch.object(self.host, 'execute') as execute:
            with self.assertRaisesRegex(RuntimeError, 'Conflicting'):
                self.host.cleanup()
            with self.assertRaisesRegex(RuntimeError, 'Conflicting'):
                self.host.applyDNS('1.1.1.1')
            execute.assert_not_called()

    def testFailedPreparationCannotClaimAnUncapturedInterface(self):
        link = {'ifname': 'utun101', 'ifindex': 78}

        with mock.patch.object(
            self.module, '_command', return_value=json.dumps([link])
        ), mock.patch.object(self.host, 'execute') as execute:
            with self.assertRaisesRegex(RuntimeError, 'identity was not captured'):
                self.host.cleanup()

            execute.assert_not_called()
            self.assertIsNone(self.host.index)

    def testIpRouteJsonFlagsAndSeparatePrefixLengthsMatchExactRules(self):
        expected = {_linuxRuleSelectors(command) for command in self.host.rules}
        records = (
            (
                '-4',
                {
                    'priority': 2000,
                    'src': 'all',
                    'dst': '198.18.0.0',
                    'dstlen': 15,
                    'table': 123456,
                },
            ),
            (
                '-4',
                {
                    'priority': 2002,
                    'src': 'all',
                    'not': None,
                    'table': 'main',
                    'suppress_prefixlength': 0,
                    'dport': 53,
                    'dport_mask': '0xffff',
                },
            ),
            (
                '-4',
                {
                    'priority': 2003,
                    'src': 'all',
                    'not': None,
                    'iif': 'lo',
                    'table': 123456,
                },
            ),
            ('-4', {'priority': 2003, 'src': '0.0.0.0', 'iif': 'lo', 'table': 123456}),
            (
                '-6',
                {
                    'priority': 2002,
                    'src': 'fd00::',
                    'srclen': 64,
                    'iif': 'lo',
                    'table': 123456,
                },
            ),
            (
                '-6',
                {
                    'priority': 2001,
                    'src': 'all',
                    'iif': 'utun101',
                    'goto': 2010,
                    'unresolved': None,
                },
            ),
            (
                '-6',
                {'priority': 2003, 'src': '0', 'srclen': 1, 'iif': 'lo', 'goto': 2010},
            ),
            ('-6', {'priority': 2010, 'src': 'all', 'nop': None}),
        )

        for flag, record in records:
            with self.subTest(rule=record):
                command = _linuxRuleDelete(flag, record, 123456, 2000, 'utun101')
                self.assertIn(_linuxRuleSelectors(command), expected)

        with self.assertRaisesRegex(RuntimeError, 'Conflicting'):
            _linuxRuleDelete(
                '-4',
                {
                    'priority': 2002,
                    'not': None,
                    'dport': 53,
                    'dport_mask': '0xff00',
                    'table': 'main',
                },
                123456,
                2000,
                'utun101',
            )

    def testCleanupUsesExactSelectorsAndRetainsRefusedCleanupForRetry(self):
        plan = SingTUNHostPlan({}, platform='Linux')
        plan._linux = self.host
        plan._activated = True
        link = {'ifname': 'utun101', 'ifindex': 77}
        self.host.index = 77
        present = True
        rule = {'priority': 2000, 'src': 'all', 'dst': '198.18.0.0/15', 'table': 123456}

        def command(arguments):
            if arguments == ['ip', '-N', '-j', 'link', 'show']:
                return json.dumps([link] if present else [])
            if arguments == ['ip', '-4', '-N', '-j', 'rule', 'show']:
                return json.dumps([rule] if present else [])
            return '[]'

        with mock.patch.object(
            self.module, '_command', side_effect=command
        ), mock.patch.object(self.host, 'restoreDNS'), mock.patch.object(
            self.host, 'execute', side_effect=RuntimeError('authorization denied')
        ):
            with self.assertRaisesRegex(RuntimeError, 'authorization denied'):
                plan.cleanup(False)

        self.assertTrue(plan._activated)
        self.assertFalse(plan._cleanupComplete)

        def release(commands, **kwargs):
            nonlocal present
            self.assertTrue(kwargs['guard'])
            self.assertEqual(commands[-1], ['ip', 'link', 'delete', 'dev', 'utun101'])
            self.assertIn('to', commands[0])
            self.assertIn('198.18.0.0/15', commands[0])
            self.assertNotIn('flush', str(commands))
            present = False

        with mock.patch.object(
            self.module, '_command', side_effect=command
        ), mock.patch.object(self.host, 'restoreDNS'), mock.patch.object(
            self.host, 'execute', side_effect=release
        ) as execute:
            plan.cleanup(False)
            plan.cleanup(True)

        execute.assert_called_once()
        self.assertFalse(plan._activated)
        self.assertTrue(plan._cleanupComplete)


if __name__ == '__main__':
    unittest.main()