shopdb-flask/tests/test_plugins/test_endpoints_credentials_migration.py

"""Migration 0023: per-PC credential tables, and the client certificate cleanup.

Two starting points matter. A database that ran the 2026-09-14 builds of
0020-0022 carries app_fingerprint and five certificate columns on assets; a
fresh one carries neither. Both must reach the same schema, twice over.
"""

import importlib.util
from pathlib import Path

from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import create_engine, inspect, text


_VERSIONS = (Path(__file__).resolve().parent.parent.parent
             / 'plugins' / 'endpoints' / 'migrations' / 'versions')
_CHAIN = ('0020_app_delta_fingerprints.py', '0021_client_certificates.py',
          '0022_client_cert_revocation.py', '0023_endpoint_credentials.py')
_CERT_COLUMNS = {'client_cert_fingerprint', 'cert_issued_at', 'cert_expires_at',
                 'cert_revoked_at', 'cert_revocation_reason'}


def _load(filename):
    path = _VERSIONS / filename
    spec = importlib.util.spec_from_file_location('endpoints_mig_' + path.stem,
                                                  str(path))
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _run(connection, step='upgrade', files=_CHAIN):
    operations = Operations(MigrationContext.configure(connection))
    operations._install_proxy()
    try:
        for filename in files:
            getattr(_load(filename), step)()
    finally:
        operations._remove_proxy()
    connection.commit()


def _base_schema(connection, legacy):
    certcolumns = ''
    if legacy:
        certcolumns = (', client_cert_fingerprint VARCHAR(64), '
                       'cert_issued_at DATETIME, cert_expires_at DATETIME, '
                       'cert_revoked_at DATETIME, '
                       'cert_revocation_reason VARCHAR(255)')
    connection.execute(text(
        'CREATE TABLE assets (assetid INTEGER PRIMARY KEY, '
        'assetnumber VARCHAR(50)' + certcolumns + ')'))
    fingerprint = ', app_fingerprint VARCHAR(64)' if legacy else ''
    connection.execute(text(
        'CREATE TABLE endpointapplications (endpointapplicationid INTEGER '
        'PRIMARY KEY, assetid INTEGER, displayname VARCHAR(255)'
        + fingerprint + ')'))
    if legacy:
        connection.execute(text(
            'CREATE INDEX idx_assets_cert_fingerprint ON assets '
            '(client_cert_fingerprint)'))
        connection.execute(text(
            'CREATE UNIQUE INDEX idx_endpointapplications_fingerprint ON '
            'endpointapplications (assetid, app_fingerprint)'))
        connection.execute(text(
            "INSERT INTO endpointapplications VALUES (1, 7, 'App', '"
            + 'a' * 64 + "')"))
    connection.commit()


def _assert_head_schema(connection):
    inspector = inspect(connection)
    assetcolumns = {c['name'] for c in inspector.get_columns('assets')}
    assert not assetcolumns & _CERT_COLUMNS
    appcolumns = {c['name'] for c in inspector.get_columns('endpointapplications')}
    assert 'appfingerprint' in appcolumns
    assert 'app_fingerprint' not in appcolumns
    indexed = {tuple(ix['column_names']): ix['unique'] for ix in
               inspector.get_indexes('endpointapplications')}
    assert indexed.get(('assetid', 'appfingerprint')) in (True, 1)

    for table in ('endpointcredentials', 'endpointenrollrequests',
                  'endpointmigrationwaves'):
        assert inspector.has_table(table), table
    credentialcolumns = {c['name'] for c in inspector.get_columns('endpointcredentials')}
    assert credentialcolumns == {
        'credentialid', 'assetid', 'credentialhash', 'credentialprefix',
        'enrolledvia', 'issuedat', 'lastusedat', 'revokedat', 'revokedby',
        'reenrollblocked', 'reenrollblockedreason'}
    requestindexes = [tuple(ix['column_names']) for ix in
                      inspector.get_indexes('endpointenrollrequests')]
    assert requestindexes.count(('hostname', 'serialnumber')) == 1


def test_upgrade_from_the_client_certificate_builds():
    engine = create_engine('sqlite://')
    with engine.connect() as connection:
        _base_schema(connection, legacy=True)
        _run(connection)
        _assert_head_schema(connection)
        # Data survives the rename.
        assert connection.execute(text(
            'SELECT appfingerprint FROM endpointapplications')).scalar() == 'a' * 64

        _run(connection)
        _assert_head_schema(connection)


def test_upgrade_on_a_fresh_database():
    engine = create_engine('sqlite://')
    with engine.connect() as connection:
        _base_schema(connection, legacy=False)
        _run(connection)
        _assert_head_schema(connection)
        _run(connection)
        _assert_head_schema(connection)


def test_retired_certificate_revisions_are_no_ops():
    engine = create_engine('sqlite://')
    with engine.connect() as connection:
        _base_schema(connection, legacy=False)
        _run(connection, files=_CHAIN[1:3])
        columns = {c['name'] for c in inspect(connection).get_columns('assets')}
        assert not columns & _CERT_COLUMNS


def test_downgrade_drops_the_credential_tables():
    engine = create_engine('sqlite://')
    with engine.connect() as connection:
        _base_schema(connection, legacy=False)
        _run(connection)
        _run(connection, step='downgrade', files=_CHAIN[3:])
        inspector = inspect(connection)
        for table in ('endpointcredentials', 'endpointenrollrequests',
                      'endpointmigrationwaves'):
            assert not inspector.has_table(table), table
        _run(connection, files=_CHAIN[3:])
        _assert_head_schema(connection)