#!/usr/bin/python3
"""Serve the Router Manager: the web application and its API, which passes calls to the router.

Run by vpn-router-manager.service. It uses python3-requests to call the router and
python3-psycopg to keep the sessions; its settings and its listening loop are in
vpnrouter_manager.

Pages:
  /api/info     the version of this package and the router it manages, as JSON
  /api/session  POST signs in through the router's /auth with {"name", "password", "code"?},
                GET answers the signed-in user, DELETE signs out
  /api/session/passkey
                POST begins a sign-in with a passkey at the router's /auth/passkey;
                /api/session/passkey/complete POST signs in through /auth/passkey/complete
  /api/updates  GET the installed and the candidate version of vpn-router
                from apt (see updates()), POST updates apt's package lists (see start_check());
                needs a session, POST a member of admins
  /api/upgrade  GET the state and the output of the upgrade of that package (see upgrade_state()),
                POST starts it (see start_upgrade()); needs a session, POST a member of admins
  /api/system-updates
                GET the packages of the system apt-get upgrade would upgrade, with their installed
                and candidate versions (see system_updates()); needs a session
  /api/system-upgrade
                GET the state and the output of the upgrade of the system's packages, POST starts
                it (see start_upgrade()); needs a session, POST a member of admins
  /api/reboot   POST reboots the system (see start_reboot()); needs a member of admins
  /api/shell    a WebSocket handshake goes to the router's /shell, and the WebSocket is relayed
                (see Handler.relay()); POST passes as every other path does
  /api/<path>   the router's web service at /<path>, with the access token of the session;
                the method, the body and the answer pass unchanged; the files of a web site,
                /api/sites/<name>/files/..., pass streamed for GET, HEAD and PUT (see Handler.transfer())
  /docs/<file>  a file of the help in DOCS_DIR, Markdown, SVG or PNG (see doc_file()); needs a session
  /<file>       the web application from WEB_DIR; a path that names no file gives index.html

A session is a row of the table sessions in PostgreSQL, found by the SHA-256 of a random id in
the HttpOnly cookie COOKIE; it holds the user's access and refresh tokens and ends with the
refresh token. An access token that expires within REFRESH_MARGIN seconds is renewed through
the router's /auth/refresh before a call; a refused renewal ends the session. /api/<path>
without a session answers 401. Sign-in calls carry the browser's address in X-Forwarded-For,
which the router reads for its rate limit when the manager is one of its trusted proxies.
Without router_url the API answers 503.

The manager first initializes, without listening: on the router, unless the settings name a
database, it creates its role, password and database when they are missing; then it connects to
the sessions database and creates its table, retrying every INIT_RETRY_SECONDS until it succeeds
and logging each failure. Then it listens. Its unit starts it only once wait-for-libpq has found
libpq in the linker cache, where psycopg looks it up.

The service needs nothing else to start and does not end on its own. Every DISCOVERY_SECONDS it
works out where to listen, bind or the interfaces listen names, and runs a server at each address
(vpnrouter_manager.serve). Every request
reads the settings when it is made.
"""

import argparse
import hashlib
import http.cookies
import json
import mimetypes
import os
import re
import secrets
import signal
import socket
import ssl
import subprocess
import sys
import threading
import time
import urllib.parse
from http.server import BaseHTTPRequestHandler

sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)))

import psycopg
import requests
import urllib3

import vpnrouter_manager
from vpnrouter_manager import DATABASE, DATABASE_SECRET, ROUTER_CONFIG_FILE

#: The address a proxy on this host calls from.
LOOPBACK = '127.0.0.1'

CONFIG_FILE = vpnrouter_manager.CONFIG_FILE
WEB_DIR = '/usr/share/vpn-router/manager/web'
DOCS_DIR = '/usr/share/vpn-router/manager/docs'

#: The files of DOCS_DIR /docs serves, by extension, and their content types.
DOC_TYPES = {'.md': 'text/markdown; charset=utf-8', '.svg': 'image/svg+xml', '.png': 'image/png'}

#: The headers of every /docs answer: no type guessing, and an SVG opened on its own runs no script.
DOC_HEADERS = (('X-Content-Type-Options', 'nosniff'), ('Content-Security-Policy', "default-src 'none'; style-src 'unsafe-inline'"))

#: The cookie that holds the session id.
COOKIE = 'vpn_router_manager_session'

#: Seconds before its expiry an access token is renewed.
REFRESH_MARGIN = 30

#: Seconds between two attempts of the initialization.
INIT_RETRY_SECONDS = 15

#: The packages /api/updates reports.
PACKAGES = ('vpn-router',)

#: Touched by apt at every successful update of its package lists.
APT_UPDATED = '/var/lib/apt/periodic/update-success-stamp'

#: The transient units of the upgrade, of the reboot and of the update of the package lists, outside the manager's unit.
UPGRADE_UNIT = 'vpn-router-manager-upgrade'
SYSTEM_UPGRADE_UNIT = 'vpn-router-manager-system-upgrade'
REBOOT_UNIT = 'vpn-router-manager-reboot'
CHECK_UNIT = 'vpn-router-manager-check'

#: Seconds the update of the package lists waits for apt's lock, which Ubuntu's apt-daily or unattended-upgrades may hold.
CHECK_LOCK_SECONDS = 300

#: The upgrade helper, and the state and the output it records of the router's package and of the system's.
UPGRADE = '/usr/lib/vpn-router/upgrade'
UPGRADE_STATE = '/var/lib/vpn-router/upgrade.json'
UPGRADE_LOG = '/var/lib/vpn-router/upgrade.log'
SYSTEM_UPGRADE_STATE = '/var/lib/vpn-router/system-upgrade.json'
SYSTEM_UPGRADE_LOG = '/var/lib/vpn-router/system-upgrade.log'

#: A package apt-get -s upgrade would upgrade: Inst <name> [<installed>] (<candidate> <release> [<architecture>]).
SIMULATED = re.compile(r'^Inst (\S+) (?:\[([^\]]*)\] )?\((\S+)')

#: The last lines of the upgrade's output /api/upgrade answers.
LOG_LINES = 200

#: Written by packages that need a reboot, with the packages in REBOOT_REQUIRED.pkgs.
REBOOT_REQUIRED = '/var/run/reboot-required'

#: Seconds between the answer of /api/reboot and the reboot.
REBOOT_DELAY = 5

SCHEMA = ('CREATE TABLE IF NOT EXISTS sessions (id_hash text PRIMARY KEY, name text NOT NULL, access_token text NOT NULL, '
          'access_expires timestamptz NOT NULL, refresh_token text NOT NULL, expires timestamptz NOT NULL)')

#: Seconds between two discoveries of where the service should listen.
DISCOVERY_SECONDS = 30

#: Seconds a call to the router may take; adding a user creates an RSA key.
ROUTER_TIMEOUT = 130

#: The largest request body passed to the router, in bytes; a CSR or a certificate chain in PEM is a few kilobytes.
BODY_LIMIT = 65536

#: The largest body of a save of the Shell page's editor, /api/shell/editor/...: a file of up to 1 MiB, as the router
#: takes, escaped in JSON.
EDITOR_BODY_LIMIT = 2 * 1024 * 1024 + BODY_LIMIT

#: The paths of a web site's files, whose GET, HEAD and PUT pass streamed (see Handler.transfer()).
SITE_FILES = re.compile(r'^/api/sites/[^/]+/files(/.*)?$')


def settings():
    return vpnrouter_manager.settings(CONFIG_FILE)


def package_version():
    try:
        result = subprocess.run(['dpkg-query', '-W', '-f=${Version}', 'vpn-router'],
                                capture_output=True, text=True, timeout=5)
    except (OSError, subprocess.TimeoutExpired):
        return None
    return result.stdout.strip() if result.returncode == 0 else None


def updates():
    """The installed and the candidate version of each of PACKAGES, from apt-cache policy, and whether an update is available.

    The candidate is the version apt would install from the sources it is configured with, as of
    its last update of the package lists, whose time checked gives (APT_UPDATED, None when it is
    not known); a version apt does not know is None. checking is whether start_check's update of
    the package lists runs. Where apt cannot be asked, available is False and the packages are empty.
    """
    checking = unit_active(CHECK_UNIT)
    try:
        result = subprocess.run(['apt-cache', 'policy', *PACKAGES], capture_output=True, text=True, timeout=30)
    except (OSError, subprocess.TimeoutExpired):
        return {'available': False, 'checked': None, 'checking': checking, 'packages': []}
    try:
        checked = time.strftime('%Y-%m-%dT%H:%M:%S%z', time.localtime(os.path.getmtime(APT_UPDATED)))
    except OSError:
        checked = None
    found = {}
    package = None
    for line in result.stdout.splitlines():
        if line and not line[0].isspace() and line.endswith(':'):
            package = found.setdefault(line[:-1], {'name': line[:-1], 'installed': None, 'candidate': None})
        elif package is not None and line.strip().startswith(('Installed:', 'Candidate:')):
            key, _, value = line.strip().partition(': ')
            package[key.lower()] = None if value == '(none)' else value
    packages = [dict(found.get(name, {'name': name, 'installed': None, 'candidate': None})) for name in PACKAGES]
    for package in packages:
        package['update'] = package['installed'] is not None and package['candidate'] not in (None, package['installed'])
    return {'available': True, 'checked': checked, 'checking': checking, 'packages': packages}


def system_updates():
    """The packages of the system apt-get upgrade would upgrade, from its simulation (apt-get -s), as updates() answers
    its packages, each with update true, sorted by name; the other keys are those of updates().

    The simulation takes no lock and changes nothing; it answers what the package lists of apt's last update give.
    """
    checking = unit_active(CHECK_UNIT)
    try:
        result = subprocess.run(['apt-get', '-s', '-q', 'upgrade'], capture_output=True, text=True, timeout=60,
                                env=dict(os.environ, LC_ALL='C'))
    except (OSError, subprocess.TimeoutExpired):
        return {'available': False, 'checked': None, 'checking': checking, 'packages': []}
    if result.returncode != 0:
        return {'available': False, 'checked': None, 'checking': checking, 'packages': []}
    try:
        checked = time.strftime('%Y-%m-%dT%H:%M:%S%z', time.localtime(os.path.getmtime(APT_UPDATED)))
    except OSError:
        checked = None
    found = {}
    for line in result.stdout.splitlines():
        match = SIMULATED.match(line)
        if match:
            found[match[1]] = {'name': match[1], 'installed': match[2], 'candidate': match[3], 'update': True}
    return {'available': True, 'checked': checked, 'checking': checking, 'packages': [found[name] for name in sorted(found)]}


def unit_active(unit):
    """Whether a unit of systemd is active or activating; False where systemctl cannot be asked."""
    try:
        result = subprocess.run(['systemctl', 'is-active', f'{unit}.service', f'{unit}.timer'],
                                capture_output=True, text=True, timeout=5)
    except (OSError, subprocess.TimeoutExpired):
        return False
    return any(line.strip() in ('active', 'activating') for line in result.stdout.splitlines())


def upgrade_state(system=False):
    """The state of the upgrade of the router's package, as GET /api/upgrade answers it, or with system of the system's
    packages, as GET /api/system-upgrade answers it.

    state is idle before the first upgrade, running, succeeded or failed as the helper recorded it, and interrupted when
    it recorded running but its unit is not active, as after a reboot during the upgrade; the other keys of the helper's
    state follow, and log holds the last LOG_LINES lines of its output. reboot_required and reboot_packages tell whether
    a package asks for a reboot, and rebooting whether one is about to happen.
    """
    state_file, log_file = (SYSTEM_UPGRADE_STATE, SYSTEM_UPGRADE_LOG) if system else (UPGRADE_STATE, UPGRADE_LOG)
    try:
        with open(state_file, encoding='utf-8') as handle:
            recorded = json.load(handle)
    except (OSError, ValueError):
        recorded = {}
    if not isinstance(recorded, dict):
        recorded = {}
    running = unit_active(SYSTEM_UPGRADE_UNIT if system else UPGRADE_UNIT)
    state = recorded.get('state', 'idle')
    if running:
        state = 'running'
    elif state == 'running':
        state = 'interrupted'
    try:
        with open(log_file, encoding='utf-8', errors='replace') as handle:
            log = handle.read().splitlines()[-LOG_LINES:]
    except OSError:
        log = []
    try:
        with open(f'{REBOOT_REQUIRED}.pkgs', encoding='utf-8') as handle:
            reboot_packages = sorted({line.strip() for line in handle if line.strip()})
    except OSError:
        reboot_packages = []
    return {
        'state': state,
        'user': recorded.get('user'),
        'started': recorded.get('started'),
        'finished': recorded.get('finished'),
        'exit_status': recorded.get('exit_status'),
        'before': recorded.get('before'),
        'after': recorded.get('after'),
        'log': log,
        'reboot_required': os.path.exists(REBOOT_REQUIRED),
        'reboot_packages': reboot_packages,
        'rebooting': unit_active(REBOOT_UNIT),
    }


def transient(*args):
    """Start a transient unit with systemd-run and the arguments; the HTTP status and the message of a refusal, or None."""
    try:
        result = subprocess.run(['systemd-run', '--quiet', '--collect', *args], capture_output=True, text=True, timeout=15)
    except (OSError, subprocess.TimeoutExpired) as exc:
        return 503, f'systemd-run cannot be run here: {exc}'
    if result.returncode != 0:
        message = result.stderr.strip() or f'systemd-run ended with status {result.returncode}'
        return (409 if 'already' in message.lower() else 500), message
    return None


def busy():
    """Whether an upgrade, of the router's package or of the system's, or a reboot runs."""
    return unit_active(UPGRADE_UNIT) or unit_active(SYSTEM_UPGRADE_UNIT) or unit_active(REBOOT_UNIT)


def start_upgrade(name, system=False):
    """Start the upgrade of the router's package in the transient unit UPGRADE_UNIT, or with system of the system's
    packages in SYSTEM_UPGRADE_UNIT, outside this service, so the restarts it makes do not end it.

    The HTTP status and the JSON document: 202 when it started, 409 while an upgrade or a reboot runs, 503 where apt
    cannot be used.
    """
    if busy():
        return 409, {'error': 'an upgrade or a reboot is already running'}
    if not os.path.exists('/usr/bin/apt-get'):
        return 503, {'error': 'apt is not available here'}
    if system:
        refused = transient(f'--unit={SYSTEM_UPGRADE_UNIT}', '--description=VPN Router system upgrade',
                            '--property=Type=exec', '--no-block', '/usr/bin/python3', UPGRADE, '--system', name)
    else:
        refused = transient(f'--unit={UPGRADE_UNIT}', '--description=VPN Router package upgrade', '--property=Type=exec',
                            '--no-block', '/usr/bin/python3', UPGRADE, name)
    if refused:
        return refused[0], {'error': refused[1]}
    return 202, {'started': True}


def start_check():
    """Update apt's package lists in the transient unit CHECK_UNIT, as root outside this service's sandbox, and touch
    APT_UPDATED when it succeeds, so updates() answers the new candidates and the time of the check.

    The HTTP status and the JSON document: 202 when it started, 409 while an update of the lists, an upgrade or a
    reboot runs, 503 where apt cannot be used.
    """
    if unit_active(CHECK_UNIT) or busy():
        return 409, {'error': 'a check for updates, an upgrade or a reboot is already running'}
    if not os.path.exists('/usr/bin/apt-get'):
        return 503, {'error': 'apt is not available here'}
    refused = transient(f'--unit={CHECK_UNIT}', '--description=VPN Router package lists update', '--property=Type=oneshot',
                        f'--property=ExecStartPost=/usr/bin/touch {APT_UPDATED}', '--no-block',
                        '/usr/bin/apt-get', '-q', '-o', f'DPkg::Lock::Timeout={CHECK_LOCK_SECONDS}', 'update')
    if refused:
        return refused[0], {'error': refused[1]}
    return 202, {'checking': True}


def start_reboot():
    """Reboot the system REBOOT_DELAY seconds from now, so the answer reaches the browser first.

    The HTTP status and the JSON document: 202 when the reboot is scheduled, 409 while an upgrade or a reboot runs.
    """
    if busy():
        return 409, {'error': 'an upgrade or a reboot is already running'}
    refused = transient(f'--unit={REBOOT_UNIT}', '--description=VPN Router reboot', f'--on-active={REBOOT_DELAY}',
                        '/usr/bin/systemctl', 'reboot')
    if refused:
        return refused[0], {'error': refused[1]}
    return 202, {'rebooting': True}


def admin(current, access, client):
    """Whether the signed-in user is a member of the router's admins, as the router's /account answers."""
    response = router_call(current, 'GET', '/account', None, {'Authorization': f'Bearer {access}', 'X-Forwarded-For': client})
    try:
        return response.ok and response.json().get('admin') is True
    except ValueError:
        return False


def error(code, message):
    return code, 'application/json', json.dumps({'error': message}).encode() + b'\n'


class Unavailable(Exception):
    """A setting the API needs is missing, or the router or the database cannot be reached; the HTTP status and message."""

    def __init__(self, code, message):
        super().__init__(message)
        self.code = code


def router_call(current, method, path, body=None, headers=None):
    """The router's answer to a call; Unavailable when it cannot be reached."""
    if not current['router_url']:
        raise Unavailable(503, f'no router is configured; set router_url in {CONFIG_FILE}')
    headers = dict(headers or {}, Accept='application/json')
    if body is not None:
        headers['Content-Type'] = 'application/json'
    try:
        return requests.request(method, current['router_url'] + path, data=body, headers=headers,
                                timeout=ROUTER_TIMEOUT, allow_redirects=False)
    except requests.RequestException as failure:
        raise Unavailable(502, f'the router at {current["router_url"]} cannot be reached: {failure}')


def answer_of(response):
    """The HTTP status, the content type and the body to pass on for an answer of the router."""
    kind = response.headers.get('Content-Type', '')
    if response.ok or kind.startswith('application/json'):
        return response.status_code, kind or 'application/json', response.content
    return error(response.status_code, f'the router answered {response.status_code}: {response.text.strip()}')


def database(current):
    """A connection to the sessions database, with its table; Unavailable when it is not configured or fails."""
    if not current['database']:
        raise Unavailable(503, f'no session database is configured; set database in {CONFIG_FILE}')
    try:
        connection = psycopg.connect(current['database'], autocommit=True)
        connection.execute(SCHEMA)
    except psycopg.Error as failure:
        raise Unavailable(503, f'the session database cannot be used: {failure}')
    return connection


def psql(sql):
    """The output of an SQL statement, read from standard input, run by psql as postgres with DATABASE as the variable v1."""
    result = subprocess.run(['runuser', '-u', 'postgres', '--', 'psql', '-X', '-q', '-t', '-A', '-v', 'ON_ERROR_STOP=1',
                             '-v', f'v1={DATABASE}'], input=sql, capture_output=True, text=True)
    if result.returncode != 0:
        raise OSError(f'psql failed: {result.stderr.strip()}')
    return result.stdout.strip()


def provision():
    """Create the role, its password file and the database on the router, each when missing.

    The password, of URL-safe characters, reaches psql on standard input, not on its command line.
    """
    if not psql("SELECT 1 FROM pg_roles WHERE rolname = :'v1'"):
        if os.path.exists(DATABASE_SECRET):
            os.unlink(DATABASE_SECRET)
        psql('CREATE ROLE :"v1" LOGIN')
    if not os.path.exists(DATABASE_SECRET):
        password = secrets.token_urlsafe(32)
        psql(f"ALTER ROLE :\"v1\" PASSWORD '{password}'")
        os.makedirs(os.path.dirname(DATABASE_SECRET), exist_ok=True)
        descriptor = os.open(DATABASE_SECRET, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(descriptor, 'w', encoding='utf-8') as handle:
            handle.write(password + '\n')
    if not psql("SELECT 1 FROM pg_database WHERE datname = :'v1'"):
        psql('CREATE DATABASE :"v1" OWNER :"v1"')


def initialize():
    """Create the sessions database on the router when it is missing and connect to it, retrying every INIT_RETRY_SECONDS."""
    reported = None
    while True:
        phase = 'reading the settings'
        try:
            configured = vpnrouter_manager.text_setting(vpnrouter_manager.read_json(CONFIG_FILE), 'database')
            if not configured and os.path.exists(ROUTER_CONFIG_FILE):
                phase = 'creating the sessions database'
                provision()
            phase = 'connecting to the sessions database'
            database(settings()).close()
        except (OSError, Unavailable) as exc:
            failure = f'{phase} failed: {exc}'
            if failure != reported:
                log(f'{failure}; retrying every {INIT_RETRY_SECONDS} seconds')
                reported = failure
            time.sleep(INIT_RETRY_SECONDS)
            continue
        log('ready')
        return


def digest(session_id):
    return hashlib.sha256(session_id.encode()).hexdigest()


def tokens_of(response):
    """The tokens of a 200 answer of /auth or /auth/refresh, or None when it holds none."""
    try:
        document = response.json()
        values = (document['access_token'], int(document['expires_in']), document['refresh_token'],
                  int(document['refresh_expires_in']))
    except (ValueError, KeyError, TypeError):
        return None
    return values if response.status_code == 200 and all(values) else None


def sign_in(current, path, body, client):
    """Sign in through the router's path, /auth or /auth/passkey/complete: the HTTP status, the content type, the body
    and the session id, or None."""
    response = router_call(current, 'POST', path, body, {'X-Forwarded-For': client})
    tokens = tokens_of(response)
    if tokens is None:
        return (*answer_of(response), None)
    try:
        name = response.json()['name']
    except (ValueError, KeyError, TypeError):
        return (*error(502, 'the router named no user'), None)
    session_id = secrets.token_urlsafe(32)
    access, access_seconds, refresh, refresh_seconds = tokens
    with database(current) as connection:
        connection.execute('DELETE FROM sessions WHERE expires < now()')
        connection.execute('INSERT INTO sessions VALUES (%s, %s, %s, now() + make_interval(secs => %s), %s, '
                           'now() + make_interval(secs => %s))',
                           (digest(session_id), name, access, access_seconds, refresh, refresh_seconds))
    return 200, 'application/json', json.dumps({'name': name}).encode() + b'\n', session_id


def session(current, session_id, client):
    """The user name and a current access token of the session, renewed when it expires soon, or None."""
    if not session_id:
        return None
    with database(current) as connection:
        with connection.transaction():
            row = connection.execute('SELECT name, access_token, access_expires < now() + make_interval(secs => %s), '
                                     'refresh_token FROM sessions WHERE id_hash = %s AND expires > now() FOR UPDATE',
                                     (REFRESH_MARGIN, digest(session_id))).fetchone()
            if row is None:
                return None
            name, access, stale, refresh = row
            if not stale:
                return name, access
            response = router_call(current, 'POST', '/auth/refresh', json.dumps({'refresh_token': refresh}).encode(),
                                   {'X-Forwarded-For': client})
            tokens = tokens_of(response)
            if tokens is None:
                connection.execute('DELETE FROM sessions WHERE id_hash = %s', (digest(session_id),))
                return None
            access, access_seconds, refresh, refresh_seconds = tokens
            connection.execute('UPDATE sessions SET access_token = %s, access_expires = now() + make_interval(secs => %s), '
                               'refresh_token = %s, expires = now() + make_interval(secs => %s) WHERE id_hash = %s',
                               (access, access_seconds, refresh, refresh_seconds, digest(session_id)))
            return name, access


def sign_out(current, session_id):
    """End the session and revoke its refresh token at the router."""
    if not session_id:
        return
    with database(current) as connection:
        row = connection.execute('DELETE FROM sessions WHERE id_hash = %s RETURNING refresh_token',
                                 (digest(session_id),)).fetchone()
    if row:
        router_call(current, 'POST', '/auth/logout', json.dumps({'refresh_token': row[0]}).encode())


def static(path):
    """A file of the web application; a path that names no file gives index.html, so the application routes it."""
    root = os.path.realpath(WEB_DIR)
    name = os.path.realpath(os.path.join(root, path.lstrip('/')))
    if not name.startswith(root + os.sep) or not os.path.isfile(name):
        name = os.path.join(root, 'index.html')
    try:
        with open(name, 'rb') as handle:
            data = handle.read()
    except OSError:
        return 404, 'text/plain; charset=utf-8', b'Not found\n'
    return 200, mimetypes.guess_type(name)[0] or 'application/octet-stream', data


def doc_file(path, root=DOCS_DIR):
    """The file of the help that a /docs/<file> path names and its content type, or None: a file inside root whose
    extension DOC_TYPES lists."""
    root = os.path.realpath(root)
    name = os.path.realpath(os.path.join(root, urllib.parse.unquote(path).lstrip('/')))
    kind = DOC_TYPES.get(os.path.splitext(name)[1].lower())
    if kind is None or not name.startswith(root + os.sep) or not os.path.isfile(name):
        return None
    return name, kind


class Stream:
    """A request body of length bytes read from a stream as requests sends it, which takes its length from len()."""

    def __init__(self, stream, length):
        self.stream = stream
        self.length = length
        self.left = length

    def __len__(self):
        return self.length

    def read(self, size=-1):
        size = self.left if size is None or size < 0 else min(size, self.left)
        data = self.stream.read(size) if size else b''
        self.left -= len(data)
        return data


class Handler(BaseHTTPRequestHandler):
    protocol_version = 'HTTP/1.1'

    def session_id(self):
        """The session id of the request's cookie, or ''."""
        cookies = http.cookies.SimpleCookie()
        try:
            cookies.load(self.headers.get('Cookie', ''))
        except http.cookies.CookieError:
            return ''
        return cookies[COOKIE].value if COOKIE in cookies else ''

    def proxied(self):
        """True for a call from a proxy on this host, such as the router's HTTPS proxy, which names the client."""
        return self.client_address[0] == LOOPBACK and 'X-Forwarded-For' in self.headers

    def client(self):
        """The browser's address: the last X-Forwarded-For entry of a proxied call, else the peer."""
        if self.proxied():
            return self.headers['X-Forwarded-For'].split(',')[-1].strip()
        return self.client_address[0]

    def api(self, path, route, payload):
        """Answer an /api/ request: the HTTP status, the content type, the body and the Set-Cookie value or None."""
        current = settings()
        client = self.client()
        session_id = self.session_id()
        if route == '/api/session/passkey' and self.command == 'POST':
            return (*answer_of(router_call(current, 'POST', '/auth/passkey', b'{}', {'X-Forwarded-For': client})), None)
        if route in ('/api/session', '/api/session/passkey/complete') and self.command == 'POST':
            code, kind, body, created = sign_in(current, '/auth' if route == '/api/session' else '/auth/passkey/complete',
                                                payload, client)
            if created is None:
                return code, kind, body, None
            return code, kind, body, f'{COOKIE}={created}; HttpOnly; SameSite=Strict; Path=/'
        if route == '/api/session' and self.command == 'DELETE':
            sign_out(current, session_id)
            return 200, 'application/json', b'{"signed_out": true}\n', f'{COOKIE}=; HttpOnly; SameSite=Strict; Path=/; Max-Age=0'
        found = session(current, session_id, client)
        if found is None:
            return (*error(401, 'not signed in'), None)
        name, access = found
        if route == '/api/session':
            return 200, 'application/json', json.dumps({'name': name}).encode() + b'\n', None
        if route == '/api/updates' and self.command in ('GET', 'HEAD'):
            return 200, 'application/json', json.dumps(updates()).encode() + b'\n', None
        if route == '/api/updates' and self.command == 'POST':
            if not admin(current, access, client):
                return (*error(403, f'user {name} is not a member of admins'), None)
            code, document = start_check()
            return code, 'application/json', json.dumps(document).encode() + b'\n', None
        if route == '/api/system-updates' and self.command in ('GET', 'HEAD'):
            return 200, 'application/json', json.dumps(system_updates()).encode() + b'\n', None
        if route in ('/api/upgrade', '/api/system-upgrade') and self.command in ('GET', 'HEAD'):
            return 200, 'application/json', json.dumps(upgrade_state(route == '/api/system-upgrade')).encode() + b'\n', None
        if route in ('/api/upgrade', '/api/system-upgrade', '/api/reboot') and self.command == 'POST':
            if not admin(current, access, client):
                return (*error(403, f'user {name} is not a member of admins'), None)
            code, document = start_reboot() if route == '/api/reboot' else start_upgrade(name, route == '/api/system-upgrade')
            return code, 'application/json', json.dumps(document).encode() + b'\n', None
        return (*answer_of(router_call(current, self.command, path[4:], payload,
                                       {'Authorization': f'Bearer {access}', 'X-Forwarded-For': client})), None)

    def relay(self, path):
        """Pass the WebSocket handshake of /api/shell to the router's /shell with the session's access token, then
        relay the bytes both ways until either side closes; the HTTP status, the content type and the body of a
        refusal, or None after a relayed WebSocket.

        The handshake must come from a page of this host (Origin), and asks the router for no extension, so the
        frames pass unchanged.
        """
        current = settings()
        client = self.client()
        if urllib.parse.urlsplit(self.headers.get('Origin', '')).netloc != self.headers.get('Host', ''):
            return error(403, 'the shell must be opened from a page of the Router Manager')
        found = session(current, self.session_id(), client)
        if found is None:
            return error(401, 'not signed in')
        if not current['router_url']:
            raise Unavailable(503, f'no router is configured; set router_url in {CONFIG_FILE}')
        url = urllib.parse.urlsplit(current['router_url'] + path[4:])
        secure = url.scheme == 'https'
        request = '\r\n'.join([
            f'GET {url.path}?{url.query} HTTP/1.1', f'Host: {url.netloc}', 'Upgrade: websocket', 'Connection: Upgrade',
            f'Sec-WebSocket-Key: {self.headers.get("Sec-WebSocket-Key", "")}',
            f'Sec-WebSocket-Version: {self.headers.get("Sec-WebSocket-Version", "")}',
            f'Authorization: Bearer {found[1]}', f'X-Forwarded-For: {client}', '', ''])
        try:
            upstream = socket.create_connection((url.hostname, url.port or (443 if secure else 80)), timeout=ROUTER_TIMEOUT)
            if secure:
                upstream = ssl.create_default_context().wrap_socket(upstream, server_hostname=url.hostname)
            upstream.sendall(request.encode())
            reader = upstream.makefile('rb')
            code = int(reader.readline().split(b' ', 2)[1])
            headers = {}
            while (line := reader.readline().strip()):
                name, _, value = line.decode('latin-1').partition(':')
                headers[name.strip().lower()] = value.strip()
            if code != 101:
                body = reader.read(int(headers.get('content-length', '0')))
                upstream.close()
                return code, headers.get('content-type', 'application/json'), body
            upstream.settimeout(None)
        except (OSError, ValueError, IndexError) as failure:
            raise Unavailable(502, f'the router at {current["router_url"]} cannot open a shell: {failure}')
        self.close_connection = True
        try:
            self.send_response(101, 'Switching Protocols')
            self.send_header('Upgrade', 'websocket')
            self.send_header('Connection', 'Upgrade')
            self.send_header('Sec-WebSocket-Accept', headers.get('sec-websocket-accept', ''))
            self.end_headers()
        except ConnectionError:
            # The client went away before the shell was relayed; the router ends the shell when upstream closes.
            upstream.close()
            return None

        def pump(source, target):
            try:
                while data := source.read1(65536):
                    target.sendall(data)
            except OSError:
                pass
            finally:
                for end in (upstream, self.connection):
                    try:
                        end.shutdown(socket.SHUT_RDWR)
                    except OSError:
                        pass

        thread = threading.Thread(target=pump, args=(reader, self.connection), name='shell output', daemon=True)
        thread.start()
        pump(self.rfile, upstream)
        thread.join()
        upstream.close()
        return None

    def transfer(self, path, send_body):
        """Pass a GET, HEAD or PUT of a web site's files to the router with the session's access token, streaming the
        request body and the answer; the HTTP status, the content type and the body of a refusal, or None once the
        router's answer was passed on.

        The body of a PUT goes to the router as it arrives, with its Content-Type and Content-Length, so it is not held
        to BODY_LIMIT; the answer passes with its Content-Type, Content-Length and Content-Disposition.
        """
        current = settings()
        client = self.client()
        found = session(current, self.session_id(), client)
        if found is None:
            return error(401, 'not signed in')
        if not current['router_url']:
            raise Unavailable(503, f'no router is configured; set router_url in {CONFIG_FILE}')
        headers = {'Authorization': f'Bearer {found[1]}', 'X-Forwarded-For': client}
        body = None
        if self.command == 'PUT':
            length = self.headers.get('Content-Length', '')
            if not length.isdigit():
                return error(411, 'the body needs a Content-Length')
            headers['Content-Type'] = self.headers.get('Content-Type') or 'application/octet-stream'
            body = Stream(self.rfile, int(length))
        try:
            response = requests.request(self.command, current['router_url'] + path[4:], data=body, headers=headers,
                                        timeout=ROUTER_TIMEOUT, allow_redirects=False, stream=True)
        except requests.RequestException as failure:
            raise Unavailable(502, f'the router at {current["router_url"]} cannot be reached: {failure}')
        with response:
            try:
                self.send_response(response.status_code)
                for name in ('Content-Type', 'Content-Length', 'Content-Disposition'):
                    if name in response.headers:
                        self.send_header(name, response.headers[name])
                self.send_header('Cache-Control', 'no-store')
                self.end_headers()
                if send_body:
                    for chunk in response.raw.stream(65536, decode_content=False):
                        self.wfile.write(chunk)
            except ConnectionError:
                # The client went away before the answer was passed on.
                self.close_connection = True
            except urllib3.exceptions.HTTPError as failure:
                # The router's answer broke off after its headers were passed on.
                self.close_connection = True
                log(f'the router\'s answer to {self.command} {path.split("?", 1)[0]} broke off: {failure}')
        return None

    def doc(self, path):
        """Answer GET /docs/<path> with a file of the help (doc_file): the HTTP status, the content type and the body; 401
        without a session, 404 for a path that names no such file."""
        try:
            if session(settings(), self.session_id(), self.client()) is None:
                return error(401, 'not signed in')
        except Unavailable as failure:
            return error(failure.code, str(failure))
        except psycopg.Error as failure:
            return error(503, f'the session database failed: {failure}')
        found = doc_file(path)
        if found is None:
            return 404, 'text/plain; charset=utf-8', b'Not found\n'
        try:
            with open(found[0], 'rb') as handle:
                return 200, found[1], handle.read()
        except OSError:
            return 404, 'text/plain; charset=utf-8', b'Not found\n'

    def answer(self, send_body):
        path = self.path.split('#', 1)[0]
        route = path.split('?', 1)[0]
        cookie = None
        # A body that is not read must not be taken for the next request.
        if self.command not in ('GET', 'HEAD'):
            self.close_connection = True
        if route == '/api/info':
            https = self.proxied() and self.headers.get('X-Forwarded-Proto') == 'https'
            status = vpnrouter_manager.status_url(CONFIG_FILE, self.headers.get('Host', ''), self.server.server_address[0], https)
            body = json.dumps({'version': package_version(), 'status_url': status or None}).encode()
            code, kind = 200, 'application/json'
        elif route == '/api/shell' and self.command == 'GET' and self.headers.get('Upgrade', '').lower() == 'websocket':
            try:
                refused = self.relay(path)
            except Unavailable as failure:
                refused = error(failure.code, str(failure))
            except psycopg.Error as failure:
                refused = error(503, f'the session database failed: {failure}')
            if refused is None:
                return
            code, kind, body = refused
        elif SITE_FILES.match(route) and self.command in ('GET', 'HEAD', 'PUT'):
            try:
                refused = self.transfer(path, send_body)
            except Unavailable as failure:
                refused = error(failure.code, str(failure))
            except psycopg.Error as failure:
                refused = error(503, f'the session database failed: {failure}')
            if refused is None:
                return
            code, kind, body = refused
        elif route.startswith('/api/'):
            limit = EDITOR_BODY_LIMIT if route.startswith('/api/shell/editor/') else BODY_LIMIT
            payload = self.request_body(limit)
            if payload is False:
                code, kind, body = error(413, f'the body is larger than {limit} bytes')
            else:
                try:
                    code, kind, body, cookie = self.api(path, route, payload)
                except Unavailable as failure:
                    code, kind, body = error(failure.code, str(failure))
                except psycopg.Error as failure:
                    code, kind, body = error(503, f'the session database failed: {failure}')
        elif route.startswith('/docs/') and self.command in ('GET', 'HEAD'):
            code, kind, body = self.doc(route[len('/docs/'):])
        elif self.command in ('GET', 'HEAD'):
            code, kind, body = static(route)
        else:
            code, kind, body = 405, 'text/plain; charset=utf-8', b'Method not allowed\n'
        try:
            self.send_response(code)
            self.send_header('Content-Type', kind)
            self.send_header('Content-Length', str(len(body)))
            self.send_header('Cache-Control', 'no-store' if route.startswith('/api/') or not route.startswith('/assets/') else 'max-age=31536000, immutable')
            if cookie:
                self.send_header('Set-Cookie', cookie)
            for name, value in DOC_HEADERS if route.startswith('/docs/') else ():
                self.send_header(name, value)
            self.end_headers()
            if send_body:
                self.wfile.write(body)
        except ConnectionError:
            # The client went away before the answer was written.
            self.close_connection = True

    def request_body(self, limit=BODY_LIMIT):
        """The request body, None when there is none, or False when it is larger than limit bytes."""
        length = self.headers.get('Content-Length')
        if not length:
            return None
        try:
            size = int(length)
        except ValueError:
            return False
        if not 0 <= size <= limit:
            return False
        return self.rfile.read(size)

    def do_GET(self):
        self.answer(True)

    def do_HEAD(self):
        self.answer(False)

    def do_POST(self):
        self.answer(True)

    def do_PUT(self):
        self.answer(True)

    def do_DELETE(self):
        self.answer(True)

    def log_message(self, format, *args):
        sys.stderr.write('%s - - [%s] %s\n' % (self.client(), self.log_date_time_string(), format % args))


def log(message):
    print(f'vpn-router-manager: {message}', file=sys.stderr, flush=True)


def main():
    global CONFIG_FILE
    parser = argparse.ArgumentParser(description='Serve the VPN Router Manager.')
    parser.add_argument('--config', default=CONFIG_FILE, help=f'the settings file (default: {CONFIG_FILE})')
    CONFIG_FILE = parser.parse_args().config
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
    initialize()
    try:
        vpnrouter_manager.serve(lambda: vpnrouter_manager.listen(CONFIG_FILE), Handler, log, DISCOVERY_SECONDS)
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == '__main__':
    sys.exit(main())