"""/shell: the ticket (shell_ticket), with the passkey proof it may take (shell_passkey), the WebSocket of the terminal
(shell), and the editor of the shell: the socket vpn-router-editor names files on (editor_listen) and /shell/editor/...
(shell_editor)."""

import base64
import fcntl
import hashlib
import json
import os
import pwd
import secrets
import select
import socket
import struct
import subprocess
import termios
import threading
import time
import urllib.parse

import ldap

import vpnrouter
import vpnrouter_passkeys
import vpnrouter_users
from vpnrouter_status import run

from .common import BODY_LIMIT, MIB, USER_STATUS, log

#: Seconds a ticket of POST /shell may open a shell.
SHELL_TICKET_SECONDS = 30

#: Seconds between the checks of an open shell's user, which end the shell once its console is refused.
SHELL_CHECK_SECONDS = 60

#: The largest WebSocket message accepted from a client, in bytes.
MESSAGE_LIMIT = 65536

#: The socket vpn-router-editor (vre) names files on to open them in the editor of a shell (see editor_listen()).
EDITOR_SOCKET = '/run/vpn-router/editor.sock'

#: The helper that reads and writes a file of the editor as the shell's user.
EDITOR_FILE = '/usr/lib/vpn-router/editor-file'

#: The largest file the editor opens, in bytes, as editor-file allows; a body of PUT /shell/editor/... may hold its
#: content escaped in JSON.
EDITOR_MAX_BYTES = MIB

EDITOR_BODY_LIMIT = 2 * EDITOR_MAX_BYTES + BODY_LIMIT

#: The most files a shell has open in its editor at once.
EDITOR_FILES_MAX = 20

#: Seconds a reader of the editor socket and a run of editor-file may take.
EDITOR_TIMEOUT = 30

#: The GUID of RFC 6455 that the Sec-WebSocket-Accept of a handshake is made with.
WEBSOCKET_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

#: WebSocket opcodes.
TEXT, BINARY, CLOSE, PING, PONG = 1, 2, 8, 9, 10

#: The unused tickets of POST /shell: ticket -> (user, monotonic time of expiry).
_shell_tickets = {}

_shell_lock = threading.Lock()

#: The last use of a shell of each user, a sign-in of POST /shell or a keystroke: user -> monotonic time.
_shell_used = {}

#: The open shells by their session, which each shell has in VPN_ROUTER_SHELL: session -> {"user", "account", "uid",
#: "send", the function that sends a WebSocket frame to its client, "files": {file id: path}}.
_shell_sessions = {}


def shell_ticket(name, body):
    """Answer POST /shell: the HTTP status and the JSON document.

    The user named by the token signs in again with {"password", "code"}, its password and TOTP code, or with
    {"passkey": {"id", "credential"}}, the answer of the browser to a proof begun by POST /shell/passkey (see
    shell_passkey) with one of its own passkeys and the user verified. A wrong password or code counts as a failed
    login; it and a passkey that is not accepted are 403. Without either, the user is not asked to sign in while its
    last use of a shell (see _shell_used) is less than web.shell_auth_timeout minutes ago. The answer is {"ticket",
    "expires_in"}: a ticket that opens one shell (GET /shell) for the same user within SHELL_TICKET_SECONDS, or a null
    ticket when the user must sign in. 403 while web.shell is false, the user's console login is not on or the user has
    no TOTP token.
    """
    if not isinstance(name, str):
        return 403, {'error': 'the shell needs a token of a signed-in user'}
    if not isinstance(body, dict):
        return 400, {'error': 'the body must be a JSON object'}
    try:
        options = vpnrouter.load_config()['web']
        if not options['shell']:
            return 403, {'error': 'the shell is disabled; vpn-config-edit set web.shell true enables it'}
        refusal = console_account(name)[1]
        if refusal:
            return 403, {'error': refusal}
        if 'passkey' in body:
            passkey = body['passkey'] if isinstance(body['passkey'], dict) else {}
            with vpnrouter_users.session() as connection:
                vpnrouter_passkeys.shell_complete(connection, name, passkey.get('id'), passkey.get('credential'))
        elif 'password' in body:
            vpnrouter_users.login(name, body.get('password'), body.get('code'))
        else:
            with _shell_lock:
                used = _shell_used.get(name)
            if used is None or time.monotonic() - used >= options['shell_auth_timeout'] * 60:
                return 200, {'ticket': None, 'expires_in': 0}
    except vpnrouter_users.Failure as failure:
        if failure.status == 5:
            return 403, {'error': 'the passkey is not accepted' if 'passkey' in body else 'the password or the code is wrong'}
        return USER_STATUS.get(failure.status, 500), {'error': str(failure)}
    except (ldap.LDAPError, OSError) as exc:
        return 500, {'error': f'the user directory failed: {exc}'}
    except SystemExit as exc:
        return 500, {'error': str(exc.code)}
    ticket = secrets.token_urlsafe(32)
    now = time.monotonic()
    with _shell_lock:
        for unused in [key for key, (_, expires) in _shell_tickets.items() if expires <= now]:
            del _shell_tickets[unused]
        _shell_tickets[ticket] = (name, now + SHELL_TICKET_SECONDS)
        _shell_used[name] = now
    return 200, {'ticket': ticket, 'expires_in': SHELL_TICKET_SECONDS}


def shell_passkey(name):
    """Answer POST /shell/passkey: begin the proof of the signed-in user with one of its own passkeys for POST /shell,
    {"id", "options"}, options those of navigator.credentials.get; the HTTP status and the JSON document. 403 while
    the user may not use a shell (console_account), 400 for a user without a passkey."""
    if not isinstance(name, str):
        return 403, {'error': 'the shell needs a token of a signed-in user'}
    try:
        if not vpnrouter.load_config()['web']['shell']:
            return 403, {'error': 'the shell is disabled; vpn-config-edit set web.shell true enables it'}
        refusal = console_account(name)[1]
        if refusal:
            return 403, {'error': refusal}
        with vpnrouter_users.session() as connection:
            return 200, vpnrouter_passkeys.shell_begin(connection, name)
    except vpnrouter_users.Failure as failure:
        return USER_STATUS.get(failure.status, 500), {'error': str(failure)}
    except (ldap.LDAPError, OSError) as exc:
        return 500, {'error': f'the user directory failed: {exc}'}
    except SystemExit as exc:
        return 500, {'error': str(exc.code)}


def console_account(name):
    """The system account the console of the user runs as and None, or None and the reason the user may not use a
    console: its console login is off, or it has no TOTP token, which the PAM login session of the shell and sudo need
    also when a passkey opens the shell."""
    with vpnrouter_users.session() as connection:
        user = vpnrouter_users.summary(vpnrouter_users.entry(connection, name))
    if not user['console'] or not user['console']['enabled']:
        return None, f'console login is not enabled for {name}; vpn-router-user console sets it'
    if not user['otp']:
        return None, f'{name} has no TOTP token; the shell needs one, enrolled on the Account page'
    return user['console']['account'], None


def take_ticket(name, ticket):
    """Whether the ticket of POST /shell is unexpired and was given to the user; it is used up either way."""
    with _shell_lock:
        holder, expires = _shell_tickets.pop(ticket, (None, 0.0))
    return holder == name and expires > time.monotonic()


def shell_command(account, unit, session):
    """The command that runs the login shell of the system user account in the transient unit, on its terminal, with
    the shell's session in VPN_ROUTER_SHELL for vpn-router-editor.

    setsid makes the terminal the controlling one of systemd-run, which passes the window size on; systemd-run starts
    the shell outside the web service's sandbox, in a PAM login session of the user.
    """
    return ['setsid', '--ctty', 'systemd-run', '--quiet', '--collect', '--pty', '--send-sighup', f'--unit={unit}',
            f'--uid={account.pw_name}', '--property=PAMName=login', f'--working-directory={account.pw_dir}',
            '--setenv=TERM=xterm-256color', f'--setenv=VPN_ROUTER_SHELL={session}', account.pw_shell or '/bin/sh', '-l']


def editor_file(account, action, path, content=b'', version=''):
    """Run editor-file as the system user account outside the sandbox: read the file, or write content to it when it
    still has version. The exit status of editor-file and its output, or its message on failure."""
    command = ['systemd-run', '--quiet', '--collect', '--pipe', f'--uid={account}', EDITOR_FILE, action, path]
    try:
        result = subprocess.run(command + ([version] if action == 'write' else []), input=content, capture_output=True,
                                timeout=EDITOR_TIMEOUT)
    except (OSError, subprocess.TimeoutExpired) as exc:
        return 1, f'editor-file cannot be run: {exc}'.encode()
    return result.returncode, result.stdout if result.returncode == 0 else result.stderr.strip() or \
        f'editor-file ended with status {result.returncode}'.encode()


def shell_editor(name, method, parts, body):
    """Answer a request of a file open in the editor of one of the user's shells: the HTTP status and the JSON document.

      GET    /shell/editor/<session>/<id>   {"path", "content", "version"}, version the SHA-256 of the content
      PUT    /shell/editor/<session>/<id>   {"content", "version"} writes the content when the file still has version,
                                            else 409; answers {"path", "version"}
      DELETE /shell/editor/<session>/<id>   closes the file

    The file is read and written as the shell's console account (editor_file). 404 for a session that is not a shell of
    the user or a file it does not have open.
    """
    if len(parts) != 4:
        return 404, {'error': f'no route {method} /{"/".join(parts)}'}
    with _shell_lock:
        session = _shell_sessions.get(parts[2])
        path = session['files'].get(parts[3]) if session and session['user'] == name else None
        account = session['account'] if path else None
    if path is None:
        return 404, {'error': 'no such file is open in a shell of yours'}
    if method == 'DELETE':
        with _shell_lock:
            session['files'].pop(parts[3], None)
        return 200, {'path': path, 'closed': True}
    if method == 'GET':
        status, output = editor_file(account, 'read', path)
        if status != 0:
            return (400 if status == 2 else 500), {'error': output.decode(errors='replace')}
        return 200, {'path': path, 'content': output.decode(), 'version': hashlib.sha256(output).hexdigest()}
    if method != 'PUT':
        return 405, {'error': f'{method} is not allowed on a file of the editor'}
    if not isinstance(body, dict) or not isinstance(body.get('content'), str) or not isinstance(body.get('version'), str):
        return 400, {'error': 'the body must be {"content", "version"}'}
    content = body['content'].encode()
    if len(content) > EDITOR_MAX_BYTES:
        return 413, {'error': f'the content is larger than {EDITOR_MAX_BYTES // 1024} KiB'}
    status, output = editor_file(account, 'write', path, content, body['version'])
    if status != 0:
        return {2: 400, 3: 409}.get(status, 500), {'error': output.decode(errors='replace')}
    log(f'editor of {name} saved {path} as {account}')
    return 200, {'path': path, 'version': hashlib.sha256(content).hexdigest()}


def editor_request(connection):
    """Answer one request of vpn-router-editor on the editor socket: {"session", "path"} on a line opens the file in the
    editor of the shell of the session, whose client is sent {"editor": "open", "session", "id", "path"} as a text
    message.

    The peer must run as the shell's console account and the path must be absolute. The answer is {"ok": true} or
    {"error"} on a line.
    """
    try:
        connection.settimeout(EDITOR_TIMEOUT)
        _, uid, _ = struct.unpack('iII', connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED,
                                                                struct.calcsize('iII')))
        request = json.loads(connection.makefile('rb').readline(4096) or b'null')
        if not isinstance(request, dict) or not isinstance(request.get('session'), str) or \
                not isinstance(request.get('path'), str):
            answer = {'error': 'the request must be {"session", "path"}'}
        elif not os.path.isabs(request['path']):
            answer = {'error': 'the path must be absolute'}
        else:
            with _shell_lock:
                session = _shell_sessions.get(request['session'])
                if session is None or session['uid'] != uid:
                    answer = {'error': 'no shell of the Router Manager has this session; vre works only in its shell'}
                elif len(session['files']) >= EDITOR_FILES_MAX:
                    answer = {'error': f'the editor has {EDITOR_FILES_MAX} files open; close one first'}
                else:
                    file_id = secrets.token_urlsafe(8)
                    session['files'][file_id] = request['path']
                    answer = None
            if answer is None:
                session['send'](TEXT, json.dumps({'editor': 'open', 'session': request['session'], 'id': file_id,
                                                  'path': request['path']}).encode())
                log(f'editor of {session["user"]} opened {request["path"]}')
                answer = {'ok': True}
        connection.sendall(json.dumps(answer).encode() + b'\n')
    except (OSError, ValueError, struct.error):
        pass
    finally:
        connection.close()


def editor_listen():
    """Listen on EDITOR_SOCKET, which any local user may connect to, answering each connection in a thread with
    editor_request; the socket of an earlier run is replaced. Does not return; a failure to listen is logged and the
    editor stays closed."""
    try:
        if os.path.lexists(EDITOR_SOCKET):
            os.unlink(EDITOR_SOCKET)
        listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        listener.bind(EDITOR_SOCKET)
        os.chmod(EDITOR_SOCKET, 0o666)
        listener.listen(16)
    except OSError as exc:
        log(f'the editor socket {EDITOR_SOCKET} cannot be opened: {exc}')
        return
    while True:
        try:
            connection, _ = listener.accept()
        except OSError:
            continue
        threading.Thread(target=editor_request, args=(connection,), name='editor request', daemon=True).start()


def frame(opcode, payload):
    """A final, unmasked WebSocket frame of a server."""
    size = len(payload)
    if size < 126:
        length = bytes([size])
    elif size < 65536:
        length = bytes([126]) + size.to_bytes(2, 'big')
    else:
        length = bytes([127]) + size.to_bytes(8, 'big')
    return bytes([0x80 | opcode]) + length + payload


def exact(handler, size):
    """The next size bytes from the client of the Handler handler; ConnectionError when it closes first."""
    data = handler.rfile.read(size)
    if len(data) != size:
        raise ConnectionError('the client closed the connection')
    return data


def message(handler, send):
    """The opcode and the payload of the next WebSocket data message of the client of the Handler handler, or CLOSE
    and its payload.

    Fragments are joined and unmasked; a ping is answered with send, a pong ignored. ValueError for an unmasked
    frame or a message over MESSAGE_LIMIT.
    """
    opcode, payload = None, b''
    while True:
        head = exact(handler, 2)
        code, size = head[0] & 0x0F, head[1] & 0x7F
        if not head[1] & 0x80:
            raise ValueError('a frame of the client is not masked')
        if size == 126:
            size = int.from_bytes(exact(handler, 2), 'big')
        elif size == 127:
            size = int.from_bytes(exact(handler, 8), 'big')
        if len(payload) + size > MESSAGE_LIMIT:
            raise ValueError(f'a message is larger than {MESSAGE_LIMIT} bytes')
        mask = exact(handler, 4) * (size // 4 + 1)
        data = bytes(byte ^ key for byte, key in zip(exact(handler, size), mask))
        if code == PING:
            send(PONG, data)
            continue
        if code == PONG:
            continue
        if code == CLOSE:
            return CLOSE, data
        opcode = code if opcode is None else opcode
        payload += data
        if head[0] & 0x80:
            return opcode, payload


def shell(handler, name):
    """Answer GET /shell?ticket=<ticket>, a WebSocket handshake the Handler handler received, with a login shell of the
    user's console account (vpnrouter_users.console) on a terminal.

    The ticket must come from POST /shell for the user named by the token. Binary messages of the client are its
    keystrokes, a text message {"cols", "rows"} sets the window size, and binary messages of the server are the
    terminal's output; a text message of the server {"editor": "open", "session", "id", "path"} opens a file the
    shell named with vpn-router-editor (editor_request), which the client reads and saves with
    /shell/editor/<session>/<id> (shell_editor) while the shell lives. The shell ends after web.shell_idle_minutes
    without a keystroke, when it exits, when the client closes the connection, or when a check of the user (SHELL_CHECK_SECONDS) finds its console login off or
    its TOTP token gone; the server then closes it with the reason. Each opening, end and refusal is logged
    with its reason.
    """
    def refuse(code, message):
        log(f'shell of {name} from {handler.client()} refused: {message}')
        body = json.dumps({'error': message}).encode() + b'\n'
        try:
            handler.send_response(code)
            handler.send_header('Content-Type', 'application/json')
            handler.send_header('Content-Length', str(len(body)))
            handler.end_headers()
            handler.wfile.write(body)
        except ConnectionError:
            handler.close_connection = True

    key = handler.headers.get('Sec-WebSocket-Key', '').strip()
    if handler.headers.get('Upgrade', '').lower() != 'websocket' or 'upgrade' not in handler.headers.get('Connection', '').lower() \
            or handler.headers.get('Sec-WebSocket-Version') != '13' or not key:
        return refuse(400, 'GET /shell must be a WebSocket handshake')
    query = urllib.parse.parse_qs(urllib.parse.urlsplit(handler.path).query)
    if not take_ticket(name, query.get('ticket', [''])[0]):
        return refuse(403, 'the ticket is not valid; POST /shell gives one')
    try:
        options = vpnrouter.load_config()['web']
    except SystemExit as exc:
        return refuse(500, str(exc.code))
    if not options['shell']:
        return refuse(403, 'the shell is disabled; vpn-config-edit set web.shell true enables it')
    try:
        login, refusal = console_account(name)
    except (vpnrouter_users.Failure, ldap.LDAPError, OSError) as exc:
        return refuse(500, f'the user directory failed: {exc}')
    if refusal:
        return refuse(403, refusal)
    try:
        account = pwd.getpwnam(login)
    except KeyError:
        return refuse(500, f'the system does not know the account {login}; is nslcd running?')
    unit = f'vpn-router-shell-{secrets.token_hex(4)}'
    session = secrets.token_urlsafe(24)
    master, terminal = os.openpty()
    try:
        process = subprocess.Popen(shell_command(account, unit, session), stdin=terminal, stdout=terminal, stderr=terminal,
                                   env={'PATH': '/usr/sbin:/usr/bin:/sbin:/bin'})
    except OSError as exc:
        os.close(master)
        return refuse(500, f'the shell cannot be started: {exc}')
    finally:
        os.close(terminal)
    accept = base64.b64encode(hashlib.sha1((key + WEBSOCKET_GUID).encode()).digest()).decode()
    handler.close_connection = True
    try:
        handler.send_response(101, 'Switching Protocols')
        handler.send_header('Upgrade', 'websocket')
        handler.send_header('Connection', 'Upgrade')
        handler.send_header('Sec-WebSocket-Accept', accept)
        handler.end_headers()
    except ConnectionError:
        # The client went away before the shell opened; the shell started for it is stopped.
        run('systemctl', 'stop', unit)
        os.close(master)
        process.kill()
        process.wait()
        log(f'shell of {name} from {handler.client()} ended before it opened: the client closed the connection')
        return
    client, started, idle = handler.client(), time.monotonic(), options['shell_idle_minutes'] * 60
    log(f'shell of {name} from {client} opened as {account.pw_name} in {unit}')

    lock, ended, typed = threading.Lock(), threading.Event(), [started]

    def send(opcode, payload):
        with lock:
            handler.wfile.write(frame(opcode, payload))

    with _shell_lock:
        _shell_sessions[session] = {'user': name, 'account': account.pw_name, 'uid': account.pw_uid, 'send': send,
                                    'files': {}}

    def receive():
        """Pass the client's keystrokes and window sizes to the terminal until it closes the connection."""
        try:
            while True:
                opcode, payload = message(handler, send)
                if opcode == CLOSE:
                    break
                if opcode == BINARY:
                    os.write(master, payload)
                    typed[0] = time.monotonic()
                    with _shell_lock:
                        _shell_used[name] = typed[0]
                elif opcode == TEXT:
                    size = json.loads(payload)
                    fcntl.ioctl(master, termios.TIOCSWINSZ, struct.pack('HHHH', int(size['rows']), int(size['cols']), 0, 0))
        except (OSError, ValueError, KeyError, TypeError, struct.error):
            pass
        finally:
            ended.set()

    reader = threading.Thread(target=receive, name=f'{unit} input', daemon=True)
    reader.start()
    reason = 'the shell ended'
    checked = started
    try:
        while not ended.is_set():
            left = typed[0] + idle - time.monotonic()
            if left <= 0:
                reason = f'no keystroke for {options["shell_idle_minutes"]} minutes'
                break
            # The user's console login and TOTP token are checked again every SHELL_CHECK_SECONDS.
            if time.monotonic() - checked >= SHELL_CHECK_SECONDS:
                checked = time.monotonic()
                try:
                    refusal = console_account(name)[1]
                except (vpnrouter_users.Failure, ldap.LDAPError, OSError) as exc:
                    refusal = f'the user directory failed: {exc}'
                if refusal:
                    reason = refusal
                    break
            if select.select([master], [], [], min(left, 1))[0]:
                try:
                    output = os.read(master, 65536)
                except OSError:
                    output = b''
                if not output:
                    break
                send(BINARY, output)
        else:
            reason = 'the client closed the connection'
    except OSError:
        reason = 'the connection failed'
    with _shell_lock:
        _shell_sessions.pop(session, None)
    run('systemctl', 'stop', unit)
    try:
        send(CLOSE, (1000).to_bytes(2, 'big') + reason.encode()[:120])
    except OSError:
        pass
    try:
        handler.connection.shutdown(socket.SHUT_RDWR)
    except OSError:
        pass
    reader.join()
    os.close(master)
    try:
        process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait()
    log(f'shell of {name} from {client} in {unit} ended after {int(time.monotonic() - started)} seconds: {reason}')