#!/usr/bin/python3
# -*- mode: python -*-
#
# This file is part of Plinth.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Configuration helper for the Tor service
"""

import argparse
import augeas
import codecs
import json
import os
import re
import socket
import time
import subprocess

from plinth import action_utils
from plinth.modules.tor.utils import get_real_apt_uri_path, iter_apt_uris, \
                                     get_augeas, is_running, is_enabled, \
                                     APT_TOR_PREFIX

SERVICE_FILE = '/etc/firewalld/services/tor-{0}.xml'
TOR_CONFIG = '/files/etc/tor/instances/plinth/torrc'
TOR_STATE_FILE = '/var/lib/tor-instances/plinth/state'
TOR_AUTH_COOKIE = '/var/run/tor-instances/plinth/control.authcookie'


def parse_arguments():
    """Return parsed command line arguments as dictionary"""
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='subcommand', help='Sub command')

    subparsers.add_parser('setup', help='Setup Tor configuration')
    subparsers.add_parser('get-status',
                          help='Get Tor status in JSON format')

    configure = subparsers.add_parser('configure', help='Configure Tor')
    configure.add_argument('--service', choices=['enable', 'disable'],
                           help='Configure Tor service')
    configure.add_argument('--relay', choices=['enable', 'disable'],
                           help='Configure relay')
    configure.add_argument('--bridge-relay', choices=['enable', 'disable'],
                           help='Configure bridge relay')
    configure.add_argument('--hidden-service', choices=['enable', 'disable'],
                           help='Configure hidden service')
    configure.add_argument('--apt-transport-tor',
                           choices=['enable', 'disable'],
                           help='Configure package download over Tor')
    configure.add_argument('--use-upstream-bridges',
                           choices=['enable', 'disable'],
                           help='Configure use of upstream bridges')
    configure.add_argument('--upstream-bridges',
                           help='Set list of upstream bridges to use')

    subparsers.add_parser('restart', help='Restart Tor')

    return parser.parse_args()


def subcommand_setup(_):
    """Setup Tor configuration after installing it."""

    # Disable default tor service. We will use tor@plinth instance
    # instead.
    _disable_apt_transport_tor()
    action_utils.service_disable('tor')

    subprocess.run(['tor-instance-create', 'plinth'], check=True)

    # Remove line starting with +SocksPort, since our augeas lens
    # doesn't handle it correctly.
    with open('/etc/tor/instances/plinth/torrc', 'r') as torrc:
        torrc_lines = torrc.readlines()
    with open('/etc/tor/instances/plinth/torrc', 'w') as torrc:
        for line in torrc_lines:
            if not line.startswith('+'):
                torrc.write(line)

    aug = augeas_load()

    aug.set(TOR_CONFIG + '/SocksPort[1]', '[::]:9050')
    aug.set(TOR_CONFIG + '/SocksPort[2]', '0.0.0.0:9050')
    aug.set(TOR_CONFIG + '/ControlPort', '9051')
    _enable_relay(relay='enable', bridge='enable', aug=aug)
    aug.set(TOR_CONFIG + '/ExitPolicy[1]', 'reject *:*')
    aug.set(TOR_CONFIG + '/ExitPolicy[2]', 'reject6 *:*')

    aug.set(TOR_CONFIG + '/VirtualAddrNetworkIPv4', '10.192.0.0/10')
    aug.set(TOR_CONFIG + '/AutomapHostsOnResolve', '1')
    aug.set(TOR_CONFIG + '/TransPort[1]', '127.0.0.1:9040')
    aug.set(TOR_CONFIG + '/TransPort[2]', '[::1]:9040')
    aug.set(TOR_CONFIG + '/DNSPort[1]', '127.0.0.1:9053')
    aug.set(TOR_CONFIG + '/DNSPort[2]', '[::1]:9053')

    aug.set(TOR_CONFIG + '/HiddenServiceDir',
            '/var/lib/tor-instances/plinth/hidden_service')
    aug.set(TOR_CONFIG + '/HiddenServicePort[1]',
            '22 127.0.0.1:22')
    aug.set(TOR_CONFIG + '/HiddenServicePort[2]',
            '80 127.0.0.1:80')
    aug.set(TOR_CONFIG + '/HiddenServicePort[3]',
            '443 127.0.0.1:443')

    aug.save()

    action_utils.service_enable('tor@plinth')
    action_utils.service_restart('tor@plinth')
    _update_ports()

    # wait until hidden service information is available
    tries = 0
    while not _get_hidden_service()['enabled']:
        tries += 1
        if tries >= 12:
            return

        time.sleep(10)


def subcommand_get_status(_):
    """Get Tor status in JSON format."""
    print(json.dumps(get_status()))


def subcommand_configure(arguments):
    """Configure Tor."""
    aug = augeas_load()

    if arguments.service == 'disable':
        _disable()

    _use_upstream_bridges(arguments.use_upstream_bridges, aug=aug)

    if arguments.use_upstream_bridges == 'enable':
        arguments.relay = 'disable'
        arguments.bridge_relay = 'disable'

    if arguments.upstream_bridges:
        _set_upstream_bridges(arguments.upstream_bridges, aug=aug)

    _enable_relay(arguments.relay, arguments.bridge_relay, aug=aug)

    if arguments.hidden_service == 'enable':
        _enable_hs(aug=aug)
    elif arguments.hidden_service == 'disable':
        _disable_hs(aug=aug)

    if arguments.service == 'enable':
        _enable()

    if arguments.apt_transport_tor == 'enable':
        _enable_apt_transport_tor()
    elif arguments.apt_transport_tor == 'disable':
        _disable_apt_transport_tor()


def subcommand_restart(_):
    """Restart Tor."""
    if is_enabled() and is_running():
        action_utils.service_restart('tor@plinth')

        aug = augeas_load()
        if aug.get(TOR_CONFIG + '/HiddenServiceDir'):
            # wait until hidden service information is available
            tries = 0
            while not _get_hidden_service()['enabled']:
                tries += 1
                if tries >= 12:
                    return
                time.sleep(10)


def get_status():
    """Return dict with Tor status."""
    aug = augeas_load()
    return {'use_upstream_bridges': _are_upstream_bridges_enabled(aug),
            'upstream_bridges': _get_upstream_bridges(aug),
            'relay_enabled': _is_relay_enabled(aug),
            'bridge_relay_enabled': _is_bridge_relay_enabled(aug),
            'ports': _get_ports(),
            'hidden_service': _get_hidden_service(aug)}


def _are_upstream_bridges_enabled(aug):
    """Return whether upstream bridges are being used."""
    use_bridges = aug.get(TOR_CONFIG + '/UseBridges')
    return use_bridges == '1'


def _get_upstream_bridges(aug):
    """Return upstream bridges separated by newlines."""
    matches = aug.match(TOR_CONFIG + '/Bridge')
    bridges = [aug.get(match) for match in matches]
    return '\n'.join(bridges)


def _is_relay_enabled(aug):
    """Return whether a relay is enabled."""
    orport = aug.get(TOR_CONFIG + '/ORPort')
    return bool(orport) and orport != '0'


def _is_bridge_relay_enabled(aug):
    """Return whether bridge relay is enabled."""
    bridge = aug.get(TOR_CONFIG + '/BridgeRelay')
    return bridge == '1'


def _get_ports():
    """Return dict mapping port names to numbers."""
    ports = {}
    try:
        ports['orport'] = _get_orport()
    except Exception:
        pass

    try:
        with open(TOR_STATE_FILE, 'r') as state_file:
            for line in state_file:
                matches = re.match(
                    r'^\s*TransportProxy\s+(\S*)\s+\S+:(\d+)\s*$', line)
                if matches:
                    ports[matches.group(1)] = matches.group(2)
    except FileNotFoundError:
        pass

    return ports


def _get_orport():
    """Return the ORPort by querying running instance."""
    cookie = open(TOR_AUTH_COOKIE, 'rb').read()
    cookie = codecs.encode(cookie, 'hex').decode()

    commands = '''AUTHENTICATE {cookie}
GETINFO net/listeners/or
QUIT
'''.format(cookie=cookie)

    tor_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tor_socket.connect(('localhost', 9051))
    tor_socket.send(commands.encode())
    response = tor_socket.recv(1024)
    tor_socket.close()

    line = response.split(b'\r\n')[1].decode()
    matches = re.match(r'.*="[^:]+:(\d+)"', line)
    return matches.group(1)


def _get_hidden_service(aug=None):
    """Return a string with configured Tor hidden service information"""
    hs_enabled = False
    hs_status = 'Ok'
    hs_hostname = None
    hs_ports = []

    if not aug:
        aug = augeas_load()

    hs_dir = aug.get(TOR_CONFIG + '/HiddenServiceDir')
    hs_port_paths = aug.match(TOR_CONFIG + '/HiddenServicePort')

    for hs_port_path in hs_port_paths:
        port_info = aug.get(hs_port_path).split()
        hs_ports.append({'virtport': port_info[0],
                         'target': port_info[1]})

    if not hs_dir:
        hs_status = 'Not Configured'
    else:
        try:
            with open(os.path.join(hs_dir, 'hostname'), 'r') as conf_file:
                hs_hostname = conf_file.read().strip()
                hs_enabled = True
        except Exception:
            hs_status = 'Not available (Run Tor at least once)'

    return {'enabled': hs_enabled, 'status': hs_status,
            'hostname': hs_hostname, 'ports': hs_ports}


def _enable():
    """Enable and start the service."""
    action_utils.service_enable('tor@plinth')
    _update_ports()


def _disable():
    """Disable and stop the service."""
    _disable_apt_transport_tor()
    action_utils.service_disable('tor@plinth')


def _use_upstream_bridges(use_upstream_bridges=None, aug=None):
    """Enable use of upstream bridges."""
    if use_upstream_bridges is None:
        return

    if not aug:
        aug = augeas_load()

    if use_upstream_bridges == 'enable':
        aug.set(TOR_CONFIG + '/UseBridges', '1')
    elif use_upstream_bridges == 'disable':
        aug.set(TOR_CONFIG + '/UseBridges', '0')

    aug.save()


def _set_upstream_bridges(upstream_bridges=None, aug=None):
    """Set list of upstream bridges."""
    if upstream_bridges is None:
        return

    if not aug:
        aug = augeas_load()

    aug.remove(TOR_CONFIG + '/Bridge')
    if upstream_bridges:
        bridges = [bridge.strip() for bridge in upstream_bridges.split('\n')]
        bridges = [bridge for bridge in bridges if bridge]
        for bridge in bridges:
            parts = [part for part in bridge.split() if part]
            bridge = ' '.join(parts)
            aug.set(TOR_CONFIG + '/Bridge[last() + 1]', bridge.strip())

    aug.set(TOR_CONFIG + '/ClientTransportPlugin',
            'obfs3,scramblesuit,obfs4 exec /usr/bin/obfs4proxy')

    aug.save()


def _enable_relay(relay=None, bridge=None, aug=None):
    """Enable Tor bridge relay."""
    if relay is None and bridge is None:
        return

    if not aug:
        aug = augeas_load()

    use_upstream_bridges = _are_upstream_bridges_enabled(aug)

    if relay == 'enable' and not use_upstream_bridges:
        aug.set(TOR_CONFIG + '/ORPort', 'auto')
    elif relay == 'disable':
        aug.remove(TOR_CONFIG + '/ORPort')

    if bridge == 'enable' and not use_upstream_bridges:
        aug.set(TOR_CONFIG + '/BridgeRelay', '1')
        aug.set(TOR_CONFIG + '/ServerTransportPlugin',
                'obfs3,obfs4 exec /usr/bin/obfs4proxy')
        aug.set(TOR_CONFIG + '/ExtORPort', 'auto')
    elif bridge == 'disable':
        aug.remove(TOR_CONFIG + '/BridgeRelay')
        aug.remove(TOR_CONFIG + '/ServerTransportPlugin')
        aug.remove(TOR_CONFIG + '/ExtORPort')

    aug.save()


def _enable_hs(aug=None):
    """Enable Tor hidden service"""
    if not aug:
        aug = augeas_load()

    if _get_hidden_service(aug)['enabled']:
        return

    aug.set(TOR_CONFIG + '/HiddenServiceDir',
            '/var/lib/tor-instances/plinth/hidden_service')
    aug.set(TOR_CONFIG + '/HiddenServicePort[1]',
            '22 127.0.0.1:22')
    aug.set(TOR_CONFIG + '/HiddenServicePort[2]',
            '80 127.0.0.1:80')
    aug.set(TOR_CONFIG + '/HiddenServicePort[3]',
            '443 127.0.0.1:443')
    aug.save()


def _disable_hs(aug=None):
    """Disable Tor hidden service"""
    if not aug:
        aug = augeas_load()

    if not _get_hidden_service(aug)['enabled']:
        return

    aug.remove(TOR_CONFIG + '/HiddenServiceDir')
    aug.remove(TOR_CONFIG + '/HiddenServicePort')
    aug.save()


def _enable_apt_transport_tor():
    """Enable package download over Tor."""
    try:
        aug = get_augeas()
    except Exception:
        # If there was an error, don't proceed
        print('Error: Unable to understand sources format.')
        exit(1)

    for uri_path in iter_apt_uris(aug):
        uri_path = get_real_apt_uri_path(aug, uri_path)
        uri = aug.get(uri_path)
        if uri.startswith('http://') or uri.startswith('https://'):
            aug.set(uri_path, APT_TOR_PREFIX + uri)

    aug.save()


def _disable_apt_transport_tor():
    """Disable package download over Tor."""
    try:
        aug = get_augeas()
    except Exception:
        # Disable what we can, so APT is not unusable.
        pass

    for uri_path in iter_apt_uris(aug):
        uri_path = get_real_apt_uri_path(aug, uri_path)
        uri = aug.get(uri_path)
        if uri.startswith(APT_TOR_PREFIX):
            aug.set(uri_path, uri[len(APT_TOR_PREFIX):])

    aug.save()


def _update_ports():
    """Update firewall service information."""
    ready = False
    tries = 0

    # port information may not be available immediately after Tor started
    while not ready:
        ports = _get_ports()
        ready = 'orport' in ports and 'obfs3' in ports and 'obfs4' in ports
        if ready:
            break

        tries += 1
        if tries >= 12:
            return

        time.sleep(10)

    lines = """<?xml version="1.0" encoding="utf-8"?>
<service>
  <short>Tor - {0}</short>
  <port protocol="tcp" port="{1}"/>
</service>
"""
    for name, number in ports.items():
        try:
            with open(SERVICE_FILE.format(name), 'w') as service_file:
                service_file.writelines(lines.format(name, number))
        except FileNotFoundError:
            return

    # XXX: We should ideally do firewalld reload instead.  However,
    # firewalld seems to fail to successfully reload sometimes.
    action_utils.service_restart('firewalld')


def augeas_load():
    """Initialize Augeas."""
    aug = augeas.Augeas(flags=augeas.Augeas.NO_LOAD +
                        augeas.Augeas.NO_MODL_AUTOLOAD)
    aug.set('/augeas/load/Tor/lens', 'Tor.lns')
    aug.set('/augeas/load/Tor/incl[last() + 1]',
            '/etc/tor/instances/plinth/torrc')
    aug.load()
    return aug


def main():
    """Parse arguments and perform all duties"""
    arguments = parse_arguments()

    subcommand = arguments.subcommand.replace('-', '_')
    subcommand_method = globals()['subcommand_' + subcommand]
    subcommand_method(arguments)


if __name__ == '__main__':
    main()