"""The tokens of the router's web service, for the web service and vpn-router-token.

Tokens are HS256 JWTs signed with WEB_KEY, a key the web service makes on its first run on the
router and that never leaves it. /auth issues an access token naming the user in sub, valid ACCESS_SECONDS, and
a refresh token valid REFRESH_SECONDS; vpn-router-token issues tokens with the admin claim and
any lifetime. A token without the admin claim must expire within MAX_LIFETIME seconds.

Every failure of the database raises DatabaseError, the only database exception callers see.

Refresh tokens are random strings kept in PostgreSQL as their SHA-256, in the database of
web.database, by default the vpn-router database on this host that provision() creates on the
web service's first run, with the password in DATABASE_SECRET. A refresh token is used once: refreshing replaces it.
The same database holds the state of the OpenID Provider (vpnrouter_oidc), whose tables and procedures OIDC_SQL
defines; the first connection of a process loads it. revoke() removes a user's refresh tokens and OpenID Provider
codes, tokens and sessions.
"""

import base64
import contextlib
import hashlib
import hmac
import json
import pathlib
import secrets
import subprocess
import time

import psycopg

WEB_KEY = pathlib.Path('/etc/vpn-router/web.key')
DATABASE_SECRET = pathlib.Path('/etc/vpn-router/database.secret')
OIDC_SQL = pathlib.Path('/usr/share/vpn-router/sql/oidc.sql')

#: The role and the database provision() creates when web.database is not set.
DATABASE = 'vpn-router'

#: The user operations of vpnrouter_users after which the user's refresh tokens are revoked.
REVOKING = ('passwd', 'disable', 'unlock', 'remove')

#: Seconds an access token of /auth and a refresh token are valid.
ACCESS_SECONDS = 300
REFRESH_SECONDS = 8 * 3600

#: The longest a token without the admin claim may be valid, from now, in seconds.
MAX_LIFETIME = 600

#: Seconds of clock difference accepted between a token and this router.
LEEWAY = 60

SCHEMA = 'CREATE TABLE IF NOT EXISTS refresh_tokens (hash text PRIMARY KEY, name text NOT NULL, expires timestamptz NOT NULL)'

#: Whether this process has loaded OIDC_SQL.
_oidc_loaded = False


class DatabaseError(Exception):
    """The database of the refresh tokens failed."""


def key():
    """The signing key, or '' while the web service has not made it."""
    try:
        return WEB_KEY.read_text(encoding='utf-8').strip()
    except OSError:
        return ''


def ensure_key():
    """Make the signing key when it is missing."""
    if not key():
        WEB_KEY.parent.mkdir(parents=True, exist_ok=True)
        WEB_KEY.touch(mode=0o600)
        WEB_KEY.chmod(0o600)
        WEB_KEY.write_text(secrets.token_urlsafe(36) + '\n', encoding='utf-8')


def b64encode(data):
    return base64.urlsafe_b64encode(data).rstrip(b'=').decode()


def b64decode(part):
    return base64.urlsafe_b64decode(part + '=' * (-len(part) % 4))


def sign(claims, secret):
    head = b64encode(json.dumps({'alg': 'HS256', 'typ': 'JWT'}).encode())
    body = b64encode(json.dumps(claims).encode())
    signature = hmac.new(secret.encode(), f'{head}.{body}'.encode(), hashlib.sha256).digest()
    return f'{head}.{body}.{b64encode(signature)}'


def verify(token, secret):
    """The claims of a token signed with HS256 by the secret, or None when it is not acceptable.

    Only HS256 is accepted, so a token cannot choose a weaker algorithm. The token must expire
    (exp); without the admin claim not later than MAX_LIFETIME seconds from now, so that a stolen
    token is of little use. LEEWAY seconds of clock difference are allowed.
    """
    try:
        header, payload, signature = token.split('.')
        if json.loads(b64decode(header)).get('alg') != 'HS256':
            return None
        expected = hmac.new(secret.encode(), f'{header}.{payload}'.encode(), hashlib.sha256).digest()
        if not hmac.compare_digest(expected, b64decode(signature)):
            return None
        claims = json.loads(b64decode(payload))
        now = time.time()
        expires = claims['exp']
        if isinstance(expires, bool) or not isinstance(expires, (int, float)) or expires < now - LEEWAY:
            return None
        if claims.get('admin') is not True and expires > now + MAX_LIFETIME + LEEWAY:
            return None
        not_before = claims.get('nbf')
        if not_before is not None and (isinstance(not_before, bool) or not isinstance(not_before, (int, float)) or not_before > now + LEEWAY):
            return None
    except (ValueError, KeyError, TypeError, AttributeError):
        return None
    return claims


def admin_token(seconds):
    """A token with the admin claim, valid for the seconds given."""
    secret = key()
    if not secret:
        raise OSError(f'{WEB_KEY} does not exist; vpn-router-web makes it when it starts')
    now = int(time.time())
    return {'token': sign({'admin': True, 'iat': now, 'exp': now + seconds}, secret), 'expires': now + seconds}


def local_connection(name, secret):
    """The connection string of a database provision_database() created; its password, of URL-safe characters, is
    written unquoted."""
    password = secret.read_text(encoding='utf-8').strip()
    return f'host=localhost dbname={name} user={name} password={password}'


def connection_string(cfg):
    """web.database, or the database provision() creates."""
    configured = cfg['web']['database'].strip()
    if configured:
        return configured
    return local_connection(DATABASE, DATABASE_SECRET)


@contextlib.contextmanager
def database(cfg):
    """A connection to the database of the refresh tokens, with their table and the OpenID Provider's, closed at the end;
    DatabaseError on any failure."""
    global _oidc_loaded
    try:
        oidc = None if _oidc_loaded else OIDC_SQL.read_text(encoding='utf-8')
    except OSError as exc:
        raise DatabaseError(f'{OIDC_SQL} cannot be read: {exc}') from exc
    try:
        with psycopg.connect(connection_string(cfg), autocommit=True) as connection:
            connection.execute(SCHEMA)
            if oidc is not None:
                connection.execute(oidc)
                _oidc_loaded = True
            yield connection
    except psycopg.Error as exc:
        raise DatabaseError(f'the token database failed: {exc}') from exc


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


def issue(cfg, name):
    """An access token and a refresh token for the user, as /auth answers them."""
    secret = key()
    now = int(time.time())
    refresh = secrets.token_urlsafe(32)
    with database(cfg) as connection:
        connection.execute('DELETE FROM refresh_tokens WHERE expires < now()')
        connection.execute("INSERT INTO refresh_tokens VALUES (%s, %s, now() + make_interval(secs => %s))",
                           (digest(refresh), name, REFRESH_SECONDS))
    return {
        'access_token': sign({'sub': name, 'iat': now, 'exp': now + ACCESS_SECONDS}, secret),
        'token_type': 'Bearer',
        'expires_in': ACCESS_SECONDS,
        'refresh_token': refresh,
        'refresh_expires_in': REFRESH_SECONDS,
    }


def take(cfg, token):
    """Remove an unexpired refresh token and return the user it was issued to, or None."""
    if not isinstance(token, str) or not token:
        return None
    with database(cfg) as connection:
        row = connection.execute('DELETE FROM refresh_tokens WHERE hash = %s RETURNING name, expires > now()',
                                 (digest(token),)).fetchone()
    return row[0] if row and row[1] else None


def revoke(cfg, name):
    """Remove all refresh tokens of the user, and its codes, tokens and sessions of the OpenID Provider."""
    with database(cfg) as connection:
        connection.execute('DELETE FROM refresh_tokens WHERE name = %s', (name,))
        connection.execute('CALL oidc_revoke(%s)', (name,))


def psql(sql, name):
    """The output of an SQL statement, read from standard input, run by psql as postgres with name as the variable v1."""
    args = ['runuser', '-u', 'postgres', '--', 'psql', '-X', '-q', '-t', '-A', '-v', 'ON_ERROR_STOP=1', '-v', f'v1={name}']
    result = subprocess.run(args, 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_database(name, secret):
    """Create the role of the name, its password file secret and its database of the same name when they are 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'", name):
        secret.unlink(missing_ok=True)
        psql("CREATE ROLE :\"v1\" LOGIN", name)
    if not secret.is_file():
        password = secrets.token_urlsafe(32)
        psql(f"ALTER ROLE :\"v1\" PASSWORD '{password}'", name)
        secret.touch(mode=0o600)
        secret.chmod(0o600)
        secret.write_text(password + '\n', encoding='utf-8')
    if not psql("SELECT 1 FROM pg_database WHERE datname = :'v1'", name):
        psql("CREATE DATABASE :\"v1\" OWNER :\"v1\"", name)


def reset_password(name, secret):
    """Give the role of the name the password of its file secret, when both exist, as after a restore."""
    if secret.is_file() and psql("SELECT 1 FROM pg_roles WHERE rolname = :'v1'", name):
        password = secret.read_text(encoding='utf-8').strip()
        psql(f"ALTER ROLE :\"v1\" PASSWORD '{password}'", name)


def provision(cfg):
    """Create the role, its password file and the database of the refresh tokens when web.database is not set and they
    are missing."""
    if not cfg['web']['database'].strip():
        provision_database(DATABASE, DATABASE_SECRET)