LorenEteval_Furious/Furious/Frozenlib/SingTUNRouting.py

# 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/>.

"""Low-level sing-tun Linux routing, DNS and privileged helper execution."""

from __future__ import annotations

from Furious.Frozenlib.Utility import runExternalCommand
from Furious.Frozenlib.SystemRoutingTable import SystemRoutingTable

import copy
import ipaddress
import json
import os
import pathlib
import secrets
import shlex
import shutil
import subprocess
import tempfile
import time

__all__ = ['SingTUNLinuxRouting']


def _command(arguments):
    """Run a checked, bounded argument vector without shell interpolation."""
    try:
        return (
            runExternalCommand(
                arguments,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
                timeout=5,
            )
            .stdout.decode('utf-8', 'replace')
            .strip()
        )
    except (OSError, subprocess.SubprocessError):
        # Command output may contain host configuration; keep the UI diagnostic
        # useful without forwarding it or profile credentials to logs.
        raise RuntimeError('sing-tun host operation failed: ' + arguments[0]) from None


def _linuxRuleDelete(flag, rule, table, start, device):
    """Match a pinned native rule's selectors, never delete a priority wildcard."""
    allowed = {
        'priority',
        'src',
        'dst',
        'table',
        'not',
        'iif',
        'oif',
        'goto',
        'suppress_prefixlength',
        'suppress_prefixlen',
        'dport',
        'action',
        'protocol',
        'srclen',
        'dstlen',
        'nop',
        'unresolved',
        'iif_detached',
        'dport_start',
        'dport_end',
        'dport_mask',
    }

    priority = int(rule['priority'])
    target = str(rule.get('table', ''))

    if set(rule) - allowed or target not in ('', str(table), 'main', '254'):
        raise RuntimeError('Conflicting Linux policy rule; refusing sing-tun recovery')

    if rule.get('iif') not in (None, 'lo', device) or rule.get('oif') is not None:
        raise RuntimeError(
            'Conflicting Linux interface rule; refusing sing-tun recovery'
        )

    if 'goto' in rule and rule['goto'] != start + 10:
        raise RuntimeError('Conflicting Linux goto rule; refusing sing-tun recovery')

    for key in ('nop', 'unresolved', 'iif_detached'):
        if key in rule and rule[key] is not None:
            raise RuntimeError('Conflicting Linux policy-rule state')

    if rule.get('action') not in (None, 'nop'):
        raise RuntimeError('Conflicting Linux policy-rule action')

    if 'unresolved' in rule and 'goto' not in rule:
        raise RuntimeError('Conflicting Linux goto rule')

    if 'iif_detached' in rule and 'iif' not in rule:
        raise RuntimeError('Conflicting Linux interface rule')

    for key in ('src', 'dst'):
        if key + 'len' in rule and (
            key not in rule or (rule[key] == 'all' and rule[key + 'len'] != 0)
        ):
            raise RuntimeError('Conflicting Linux rule prefix length')

    nop = rule.get('action') == 'nop' or 'nop' in rule

    if priority == start + 10 and not nop:
        raise RuntimeError(
            'Conflicting Linux terminal rule; refusing sing-tun recovery'
        )

    if rule.get('protocol') not in (None, 0, 'unspec'):
        raise RuntimeError(
            'Conflicting Linux rule protocol; refusing sing-tun recovery'
        )

    args = ['ip', flag, 'rule', 'del', 'priority', str(priority)]

    if 'not' in rule and rule['not'] is not None and type(rule['not']) is not bool:
        raise RuntimeError('Conflicting Linux inverted rule')

    if 'not' in rule and rule['not'] is not False:
        args.append('not')

    for key, parameter in (('src', 'from'), ('dst', 'to'), ('iif', 'iif')):
        if key in rule:
            value = str(rule[key])

            if key != 'iif' and value != 'all':
                if value == '0':
                    value = '0.0.0.0' if flag == '-4' else '::'

                network = ipaddress.ip_network(value, strict=False)
                length = rule.get(key + 'len', network.prefixlen)

                if type(length) is not int or not 0 <= length <= network.max_prefixlen:
                    raise RuntimeError('Conflicting Linux rule prefix length')

                if '/' in value and network.prefixlen != length:
                    raise RuntimeError('Conflicting Linux rule prefix')

                value = str(
                    ipaddress.ip_network(
                        (network.network_address, length), strict=False
                    )
                )

            args.extend([parameter, value])

    suppression = [
        rule[key]
        for key in ('suppress_prefixlen', 'suppress_prefixlength')
        if key in rule
    ]

    if suppression:
        if any(type(value) is not int or value != 0 for value in suppression):
            raise RuntimeError('Conflicting Linux suppression rule')

        args.extend(['suppress_prefixlength', '0'])

    if 'dport' in rule or 'dport_start' in rule or 'dport_end' in rule:
        port = str(rule.get('dport', ''))

        if 'dport_start' in rule or 'dport_end' in rule:
            if rule.get('dport_start') != 53 or rule.get('dport_end') != 53:
                raise RuntimeError('Conflicting Linux DNS rule')

            port = '53'

        mask = rule.get('dport_mask', 65535)

        if isinstance(mask, str):
            try:
                mask = int(mask, 0)
            except ValueError:
                raise RuntimeError('Conflicting Linux DNS rule mask') from None

        if port not in ('53', '53-53') or type(mask) is not int or mask != 65535:
            raise RuntimeError('Conflicting Linux DNS rule')

        args.extend(['dport', '53'])
    elif 'dport_mask' in rule:
        raise RuntimeError('Conflicting Linux DNS rule mask')

    if target:
        args.extend(['table', target])

    if 'goto' in rule:
        args.extend(['goto', str(rule['goto'])])
    elif nop:
        args.append('nop')
    elif not target:
        raise RuntimeError('Unrecognized Linux rule; cleanup ownership retained')

    return args


def _linuxRouteRanges(tun):
    """Build the supported Linux route ranges before handing off the device."""
    result = []

    for family in (4, 6):
        key = 'Inet' + str(family)

        if not tun[key + 'Address']:
            continue

        ranges = [
            ipaddress.ip_network(value, strict=False)
            for value in tun.get(key + 'RouteAddress')
            or ['0.0.0.0/0' if family == 4 else '::/0']
        ]

        for value in tun.get(key + 'RouteExcludeAddress', []):
            exclusion = ipaddress.ip_network(value, strict=False)
            ranges = [
                remaining
                for network in ranges
                if not network.subnet_of(exclusion)
                for remaining in (
                    network.address_exclude(exclusion)
                    if exclusion.subnet_of(network)
                    else [network]
                )
            ]

        result.extend(ipaddress.collapse_addresses(ranges))

    return result


def _linuxRules(tun, table, start):
    """Mirror the pinned Linux policy for Furious's supported native options."""
    result = []

    for family in (4, 6):
        addresses = tun['Inet' + str(family) + 'Address']

        if not addresses:
            continue

        flag = '-' + str(family)
        priority = start

        def add(*selectors):
            result.append(
                ['ip', flag, 'rule', 'add', 'priority', str(priority), *selectors]
            )

        networks = [str(ipaddress.ip_interface(value).network) for value in addresses]

        if family == 4:
            for network in networks:
                add('to', network, 'table', str(table))

            priority += 1

        add('table', str(table), 'suppress_prefixlength', '0')
        priority += 1
        add('not', 'dport', '53', 'table', 'main', 'suppress_prefixlength', '0')
        add('iif', tun['Name'], 'goto', str(start + 10))
        priority += 1

        if family == 4:
            add('not', 'iif', 'lo', 'table', str(table))
            add('iif', 'lo', 'from', '0.0.0.0/32', 'table', str(table))

        for network in networks:
            add('iif', 'lo', 'from', network, 'table', str(table))

        if family == 6:
            priority += 1
            add('iif', 'lo', 'from', '::/1', 'goto', str(start + 10))
            add('iif', 'lo', 'from', '8000::/1', 'goto', str(start + 10))
            priority += 1
            add('table', str(table))

        priority = start + 10
        add('nop')

    return result


def _linuxRuleSelectors(command):
    """Compare selector pairs rather than the ordering of ip's JSON fields."""
    selectors = []
    tokens = iter(command[6:])

    for key in tokens:
        if key in ('not', 'nop'):
            selectors.append((key, ''))
        else:
            value = next(tokens)

            if key in ('from', 'to') and value == 'all':
                continue

            if key in ('from', 'to'):
                value = str(ipaddress.ip_network(value, strict=False))

            if key == 'table' and value == '254':
                value = 'main'

            selectors.append((key, value))

    return command[1], command[5], tuple(sorted(selectors))


class SingTUNLinuxRouting:
    """Own the helper process and the persistent TUN created for one attempt."""

    @staticmethod
    def checkEnvironment():
        """Check Linux tools and the resolver before host preparation."""
        if not shutil.which('resolvectl'):
            raise RuntimeError(
                'sing-tun DNS management requires systemd-resolved on Linux; select tun2socks on other resolver configurations'
            )

        try:
            _command(['resolvectl', 'status'])
        except RuntimeError:
            raise RuntimeError(
                'sing-tun requires an active systemd-resolved service on Linux'
            ) from None

        for tool in ('ip', 'bash', *(['pkexec'] if os.geteuid() != 0 else [])):
            if not shutil.which(tool):
                raise RuntimeError('sing-tun Linux host preparation requires ' + tool)

    @staticmethod
    def reserveIdentifiers():
        """Select an unused routing table and policy-rule priority range."""
        rules = sum(
            (
                json.loads(_command(['ip', flag, '-N', '-j', 'rule', 'show']))
                for flag in ('-4', '-6')
            ),
            [],
        )

        routes = sum(
            (
                json.loads(
                    _command(['ip', flag, '-N', '-j', 'route', 'show', 'table', 'all'])
                )
                for flag in ('-4', '-6')
            ),
            [],
        )

        tables = {str(item.get('table')) for item in rules + routes}
        priorities = {int(item['priority']) for item in rules}

        for _ in range(100):
            table = secrets.randbelow(1000000000) + 100000
            start = secrets.randbelow(1800) * 11 + 1000

            if str(table) not in tables and not any(
                start <= value <= start + 10 for value in priorities
            ):
                return table, start
        else:
            raise RuntimeError('No isolated sing-tun routing identifiers are available')

    def __init__(self, tun, table, start, cancel):
        self.tun = copy.deepcopy(tun)
        self.table, self.start = table, start
        self.cancelled = cancel
        self.index = None
        self.dnsTouched = False
        self.rules = _linuxRules(tun, table, start)
        self.ranges = _linuxRouteRanges(tun)
        self._process = None
        self._directory = None
        self._holding = False
        self._ready = False

    def cancel(self):
        """Prevent a pending authorization from acquiring more host resources."""
        directory = self._directory

        if directory is not None:
            try:
                pathlib.Path(directory.name, 'cancelled').touch()
            except FileNotFoundError:
                # The worker already reaped the helper and removed its directory.
                pass

        if self._holding and not self._ready:
            self._releaseHelper()

    def _releaseHelper(self):
        """EOF requests the fixed cleanup script, never an arbitrary command."""
        if self._holding:
            self._holding = False

            if self._process is not None and self._process.stdin is not None:
                self._process.stdin.close()

    def runtimeError(self):
        """Report loss of an authorized helper while forwarding is active."""
        if self._holding and self._ready and self._process.poll() is not None:
            return 'sing-tun privileged helper exited unexpectedly'

        return ''

    def drain(self, *, wait=True):
        """Retain an unfinished helper and its script until all commands exit."""
        process = self._process

        if process is None:
            if self._directory is not None:
                self._directory.cleanup()
                self._directory = None

            return True

        if self._holding and self._ready and process.poll() is None:
            # The authorized helper stays with the live forwarding runtime.
            return True

        if process.poll() is None:
            if not wait:
                return False

            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                raise RuntimeError(
                    'sing-tun privileged helper is still running'
                ) from None

        self._process = None
        self._directory.cleanup()
        self._directory = None
        self._holding = False
        self._ready = False

        return True

    def execute(self, commands, *, preparing=False, guard=False, holding=False):
        """Run only argument vectors generated for this attempt, outside Qt."""
        self.drain()

        if self._process is not None:
            raise RuntimeError('sing-tun privileged helper is already active')

        if preparing and self.cancelled.is_set():
            raise RuntimeError('sing-tun host preparation was cancelled')

        directory = tempfile.TemporaryDirectory()
        self._directory = directory

        if preparing and self.cancelled.is_set():
            self.cancel()

        lines = ['set -eu']

        if holding:
            lines.extend(self._sessionCleanupLines())

        if guard:
            name = self.tun['Name']
            index = shlex.quote('/sys/class/net/' + name + '/ifindex')
            lines.append(f'[ "$(cat {index})" = {self.index} ]')

        for arguments in commands:
            if preparing:
                lines.append(
                    '[ -d '
                    + shlex.quote(directory.name)
                    + ' ] && [ ! -e '
                    + shlex.quote(str(pathlib.Path(directory.name, 'cancelled')))
                    + ' ]'
                )

            lines.append(' '.join(shlex.quote(str(value)) for value in arguments))

            if preparing and arguments[:3] == ['ip', 'tuntap', 'add']:
                name = self.tun['Name']
                index = shlex.quote('/sys/class/net/' + name + '/ifindex')
                lines.append(f'createdIndex=$(cat {index})')

                if not holding:
                    # Roll back device creation if a later setup command fails.
                    rollback = (
                        'code=$?; if [ "$code" -ne 0 ] '
                        f'&& [ "$(cat {index} 2>/dev/null)" = "$createdIndex" ]; then '
                        'ip link delete dev ' + shlex.quote(name) + '; fi'
                    )
                    lines.append('trap ' + shlex.quote(rollback) + ' EXIT')

            if holding and arguments[2:4] == ['rule', 'add']:
                lines.append('rulesAdded=$((rulesAdded + 1))')

        if holding:
            ready = pathlib.Path(directory.name, 'ready')
            staging = pathlib.Path(directory.name, 'ready.tmp')

            lines.extend(
                [
                    'printf "%s\\n" "$createdIndex" > ' + shlex.quote(str(staging)),
                    'mv -- '
                    + shlex.quote(str(staging))
                    + ' '
                    + shlex.quote(str(ready)),
                    # The pipe also closes if Furious crashes. No second
                    # authorization is needed to restore this attempt's state.
                    'read -r release || :',
                ]
            )

        with tempfile.NamedTemporaryFile(
            mode='w',
            encoding='utf-8',
            suffix='.sh',
            dir=directory.name,
            delete=False,
        ) as file:
            file.write('\n'.join(lines) + '\n')
            script = pathlib.Path(file.name)

        try:
            self._process = SystemRoutingTable.LinuxStartPrivilegedScript(
                str(script),
                stdin=subprocess.PIPE if holding else subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )

            self._holding = holding

            if holding:
                ready = pathlib.Path(directory.name, 'ready')

                while self._process.poll() is None:
                    if self.cancelled.is_set():
                        self.cancel()
                        raise RuntimeError('sing-tun host preparation was cancelled')

                    if ready.exists():
                        self.index = int(ready.read_text(encoding='utf-8').strip())
                        self._ready = True
                        return

                    time.sleep(0.05)

                self._releaseHelper()
                code = self._process.wait()
            else:
                # Match tun2socks: a PolicyKit interaction has no deadline.
                code = self._process.wait()
        except OSError:
            if self._process is None:
                directory.cleanup()
                self._directory = None

            raise RuntimeError(
                'Unable to start the sing-tun privileged helper'
            ) from None

        self.drain()

        if code != 0:
            raise RuntimeError(
                'sing-tun privileged host operation failed or authorization was cancelled'
            )

    def _sessionCleanupLines(self):
        """Pair successful additions with exact deletions in the same helper."""
        name = shlex.quote(self.tun['Name'])
        index = shlex.quote('/sys/class/net/' + self.tun['Name'] + '/ifindex')

        lines = [
            'createdIndex=; rulesAdded=0',
            'cleanup() {',
            'code=$?; trap - EXIT; set +e; failed=0',
        ]

        for number, rule in reversed(list(enumerate(self.rules, 1))):
            deletion = rule[:]
            deletion[3] = 'del'
            command = ' '.join(shlex.quote(str(value)) for value in deletion)
            lines.append(
                f'if [ "$rulesAdded" -ge {number} ]; then {command} || failed=1; fi'
            )

        lines.extend(
            [
                'if [ -n "$createdIndex" ] && [ -e ' + index + ' ]; then',
                f'if [ "$(cat {index})" = "$createdIndex" ]; then',
                # resolved discards per-link DNS when the link disappears,
                # including when DNS application failed partway through.
                f'resolvectl revert {name} || failed=1',
                f'ip link delete dev {name} || failed=1',
                'else failed=1; fi',
                'fi',
                '[ "$failed" -eq 0 ] || code=1',
                'exit "$code"',
                '}',
                'trap cleanup EXIT',
                "trap 'exit 1' HUP INT TERM",
            ]
        )

        return lines

    def prepare(self, dnsAddress='1.1.1.1'):
        name = self.tun['Name']

        # Recheck inside the privileged script before asking ip to create a TUN.
        commands = [['test', '!', '-e', '/sys/class/net/' + name]]
        commands.extend(
            [
                [
                    'ip',
                    'tuntap',
                    'add',
                    'dev',
                    name,
                    'mode',
                    'tun',
                    'user',
                    str(os.getuid()),
                ],
                ['ip', 'link', 'set', 'dev', name, 'mtu', str(self.tun['MTU'])],
            ]
        )

        for family in (4, 6):
            for address in self.tun['Inet' + str(family) + 'Address']:
                commands.append(
                    ['ip', '-' + str(family), 'addr', 'add', address, 'dev', name]
                )

        commands.append(['ip', 'link', 'set', 'dev', name, 'up'])

        for network in self.ranges:
            commands.append(
                [
                    'ip',
                    '-' + str(network.version),
                    'route',
                    'add',
                    str(network),
                    'dev',
                    name,
                    'table',
                    str(self.table),
                ]
            )

        commands.extend(self.rules)
        commands.extend(
            [
                ['resolvectl', 'dns', name, dnsAddress],
                ['resolvectl', 'domain', name, '~.'],
            ]
        )

        self.execute(commands, preparing=True, holding=True)
        self.dnsTouched = True

        link = self.link()

        if link is None:
            raise RuntimeError('sing-tun helper did not create the requested TUN')

        self.index = link['ifindex']

    def link(self):
        links = json.loads(_command(['ip', '-N', '-j', 'link', 'show']))
        link = next(
            (item for item in links if item['ifname'] == self.tun['Name']), None
        )

        if (
            link is not None
            and self.index is not None
            and link['ifindex'] != self.index
        ):
            raise RuntimeError(
                'Conflicting sing-tun interface; cleanup ownership retained'
            )

        if link is not None:
            if type(link.get('ifindex')) is not int or link['ifindex'] <= 0:
                raise RuntimeError('Invalid sing-tun interface identity')

        return link

    def applyDNS(self, address):
        if self.link() is None:
            raise RuntimeError('sing-tun interface disappeared before DNS confirmation')

        if self._holding and self._ready:
            # Linux DNS was configured alongside the interface/routes under
            # the original authorization. Native readiness only confirms it.
            return

        self.dnsTouched = True

        self.execute(
            [
                ['resolvectl', 'dns', self.tun['Name'], address],
                ['resolvectl', 'domain', self.tun['Name'], '~.'],
            ],
            preparing=True,
            guard=True,
        )

    def restoreDNS(self):
        if not self.dnsTouched:
            return

        if self._process is not None:
            # Do not release host routes while the forwarding child is alive.
            # cleanup() closes the helper pipe after the child is reaped.
            raise RuntimeError(
                'sing-tun DNS restoration requires stopping the forwarding child'
            )

        if self.link() is not None:
            self.execute([['resolvectl', 'revert', self.tun['Name']]], guard=True)

        self.dnsTouched = False

    def cleanup(self):
        self._releaseHelper()
        self.drain()
        link = self.link()

        if link is not None and self.index is None:
            raise RuntimeError(
                'sing-tun interface identity was not captured; cleanup ownership retained'
            )

        commands = []
        expected = set()

        if self.dnsTouched and link is not None:
            commands.append(['resolvectl', 'revert', self.tun['Name']])

        for rule in self.rules:
            deletion = rule[:]
            deletion[3] = 'del'
            expected.add(_linuxRuleSelectors(deletion))

        for flag in ('-4', '-6'):
            rules = json.loads(_command(['ip', flag, '-N', '-j', 'rule', 'show']))

            for rule in rules:
                if self.start <= int(rule['priority']) <= self.start + 10:
                    deletion = _linuxRuleDelete(
                        flag, rule, self.table, self.start, self.tun['Name']
                    )

                    if _linuxRuleSelectors(deletion) not in expected:
                        raise RuntimeError(
                            'Conflicting Linux policy rule; cleanup ownership retained'
                        )

                    commands.append(deletion)

        if link is not None:
            commands.append(['ip', 'link', 'delete', 'dev', self.tun['Name']])

        if commands:
            self.execute(commands, guard=link is not None)

        if self.link() is not None:
            raise RuntimeError('sing-tun interface remains; cleanup ownership retained')

        self.dnsTouched = False

        for flag in ('-4', '-6'):
            rules = json.loads(_command(['ip', flag, '-N', '-j', 'rule', 'show']))

            routes = [
                route
                for route in json.loads(
                    _command(['ip', flag, '-N', '-j', 'route', 'show', 'table', 'all'])
                )
                if str(route.get('table')) == str(self.table)
            ]

            if routes or any(
                self.start <= int(rule['priority']) <= self.start + 10 for rule in rules
            ):
                raise RuntimeError(
                    'sing-tun Linux routes or rules remain; cleanup ownership retained'
                )