shopdb-flask/tests/test_plugins/test_endpoints_enrollment_migration.py

"""Moving the existing fleet onto per-PC credentials.

A PC on agent 1.19 or older holds only the shared key. It moves by receiving
agent 1.20, which enrolls itself. The migration surface computes, per computer,
whether that upgrade can be pushed, and pushes it in waves through the same
enqueue path as the single-asset agentupgrade route, packaging once per wave.

The agent store is isolated per test, as in test_endpoints_agent_upgrade_route.
"""

from datetime import datetime, timedelta, timezone

import pytest

from plugins.endpoints.api import agentregistry, agentstore, agentupgradescript

MIGRATION_URL = '/api/endpoints/enrollment/migration'
WAVE_URL = '/api/endpoints/enrollment/migration/wave'

_PE_BYTES = b'MZ' + b'\x00' * 4096


@pytest.fixture(autouse=True)
def isolate_store(tmp_path, monkeypatch):
    siteagent = tmp_path / 'siteagent'
    packages = tmp_path / 'applicationpackages'
    scripts = tmp_path / 'scripts'
    for directory in (siteagent, packages, scripts):
        directory.mkdir()
    monkeypatch.setattr(agentstore, '_siteagent_dir', lambda: str(siteagent))
    monkeypatch.setattr(agentregistry, 'application_package_location',
                        lambda: (str(packages), '/api/applications/package/'))
    monkeypatch.setattr(agentupgradescript, 'script_body_location',
                        lambda: (str(scripts), '/api/scripts/body/'))
    return tmp_path


def _now():
    return datetime.now(timezone.utc).replace(tzinfo=None)


def _publish_agent(version='1.20.0'):
    agentstore.store_version('agent', version, _PE_BYTES, make_current=True)


def _pc(db, hostname, agentversion='1.19.0', serialnumber='SER',
        reportedago=timedelta(hours=1), reports=True):
    from shopdb.api import Asset, AssetType
    from plugins.computers.models import Computer
    from plugins.endpoints.models import EndpointState

    atype = AssetType.query.filter_by(assettype='computer').first()
    if atype is None:
        atype = AssetType(assettype='computer')
        db.session.add(atype)
        db.session.flush()
    asset = Asset(assetnumber=hostname + '-asset', name=hostname,
                  assettypeid=atype.assettypeid, isactive=True,
                  serialnumber=serialnumber)
    db.session.add(asset)
    db.session.flush()
    db.session.add(Computer(assetid=asset.assetid, hostname=hostname))
    if reports:
        db.session.add(EndpointState(
            assetid=asset.assetid, hostname=hostname, agentversion=agentversion,
            lastreportedat=_now() - reportedago if reportedago else None))
    db.session.commit()
    return asset.assetid


def _rows(client, headers, query=''):
    response = client.get(MIGRATION_URL + query, headers=headers)
    assert response.status_code == 200, response.get_data(as_text=True)
    return response.get_json()['data']


def _by_host(data):
    return {row['hostname']: row for row in data['computers']}


def test_readiness_states(client, db, app, auth_headers):
    from plugins.endpoints.api.agentauth import generate_credential
    from plugins.endpoints.api.enrollment import issue_credential
    from plugins.endpoints.models import EndpointCommand, EndpointEnrollRequest

    _publish_agent()
    enrolled = _pc(db, 'PC-ENROLLED')
    issue_credential(enrolled, 'auto')
    pending = _pc(db, 'PC-PENDING')
    db.session.add(EndpointEnrollRequest(
        hostname='pc-pending', serialnumber='SER', enrollnonce=generate_credential(),
        status='pending', reason='awaiting approval', matchedassetid=pending,
        requestedat=_now(), lastseenat=_now()))
    upgrading = _pc(db, 'PC-UPGRADING')
    _pc(db, 'PC-READY')
    _pc(db, 'PC-BLOCKED', serialnumber=None, reportedago=timedelta(days=3))
    _pc(db, 'PC-SHAREONLY', reports=False)
    _pc(db, 'PC-NEW', agentversion='1.20.0.0')
    db.session.commit()

    seeded_app = agentregistry.register_current_agent_package()
    assert seeded_app is not None
    scriptid = agentupgradescript.seed_upgrade_script()['scriptid']
    db.session.add(EndpointCommand(assetid=upgrading, action='run_script',
                                   target=f'script:{scriptid}',
                                   status=EndpointCommand.STATUS_PENDING))
    db.session.commit()

    data = _rows(client, auth_headers)
    rows = _by_host(data)
    assert data['target'] == '1.20.0'
    assert rows['PC-ENROLLED']['state'] == 'enrolled'
    assert rows['PC-PENDING']['state'] == 'pending'
    assert rows['PC-UPGRADING']['state'] == 'upgrading'
    assert rows['PC-READY']['state'] == 'ready'
    assert rows['PC-READY']['blockers'] == []
    assert rows['PC-BLOCKED']['state'] == 'blocked'
    assert rows['PC-BLOCKED']['blockers'] == ['offline', 'no serial on record']
    assert rows['PC-SHAREONLY']['state'] == 'notreporting'
    assert rows['PC-NEW']['state'] == 'blocked'
    assert rows['PC-NEW']['blockers'] == ['awaiting enrollment']
    assert data['counts'] == {'enrolled': 1, 'pending': 1, 'upgrading': 1,
                              'ready': 1, 'blocked': 2, 'notreporting': 1}

    only = _rows(client, auth_headers, '?state=ready')
    assert [row['hostname'] for row in only['computers']] == ['PC-READY']
    assert client.get(MIGRATION_URL + '?state=bogus',
                      headers=auth_headers).status_code == 400


def test_every_blocker_is_listed(client, db, app, auth_headers, monkeypatch):
    from plugins.endpoints import commands
    from plugins.endpoints.api import enrollment

    # No package in the store at all.
    _pc(db, 'DUP')
    _pc(db, 'dup', agentversion='1.5.0')
    lockedid = _pc(db, 'LOCKED')
    commands.PLACE_ENDPOINT_LOCK(assetid=lockedid, reason='bench repair')
    db.session.commit()
    monkeypatch.setattr(enrollment, 'minimum_agent_version_for',
                        lambda action: '1.10.0')

    rows = {}
    for row in _rows(client, auth_headers)['computers']:
        rows.setdefault(row['hostname'], []).append(row)
    assert rows['dup'][0]['blockers'] == ['duplicate hostname',
                                          'no agent package in store',
                                          'agent too old for remote upgrade']
    assert rows['LOCKED'][0]['blockers'] == ['commands locked',
                                             'no agent package in store']


def test_an_old_store_build_counts_as_no_package(client, db, app, auth_headers):
    _publish_agent('1.19.5')
    _pc(db, 'PC-A')
    row = _by_host(_rows(client, auth_headers))['PC-A']
    assert row['blockers'] == ['no agent package in store']


def test_version_compare_is_numeric(client, db, app, auth_headers):
    _publish_agent('1.20.0')
    _pc(db, 'PC-NINE', agentversion='1.9.0')
    _pc(db, 'PC-HUNDRED', agentversion='1.100.0')
    rows = _by_host(_rows(client, auth_headers))
    assert rows['PC-NINE']['state'] == 'ready'
    assert rows['PC-HUNDRED']['blockers'] == ['awaiting enrollment']


def test_wave_next_n_queues_ready_pcs_in_hostname_order(client, db, app,
                                                       auth_headers):
    from plugins.endpoints.models import EndpointCommand, EndpointMigrationWave

    _publish_agent()
    third = _pc(db, 'PC-C')
    first = _pc(db, 'PC-A')
    second = _pc(db, 'PC-B')
    _pc(db, 'PC-OFFLINE', reportedago=timedelta(days=2))

    response = client.post(WAVE_URL, json={'next': 2}, headers=auth_headers)
    assert response.status_code == 201, response.get_data(as_text=True)
    data = response.get_json()['data']
    assert data['queued'] == [first, second]
    assert data['skipped'] == []

    commands = EndpointCommand.query.filter_by(action='run_script').all()
    assert sorted(c.assetid for c in commands) == [first, second]
    wave = EndpointMigrationWave.query.one()
    assert wave.waveid == data['waveid']
    assert wave.assetidlist == [first, second]

    migration = _rows(client, auth_headers)
    rows = _by_host(migration)
    assert rows['PC-A']['state'] == 'upgrading'
    assert rows['PC-C']['state'] == 'ready'
    assert migration['waves'] == [{
        'waveid': wave.waveid, 'createdat': migration['waves'][0]['createdat'],
        'createdby': 'testadmin', 'size': 2, 'queued': 2, 'enrolled': 0,
        'failed': 0, 'done': False}]
    assert third not in data['queued']


def test_wave_explicit_ids_rechecks_readiness(client, db, app, auth_headers):
    _publish_agent()
    ready = _pc(db, 'PC-READY')
    offline = _pc(db, 'PC-OFFLINE', reportedago=timedelta(days=2))
    noagent = _pc(db, 'PC-NOAGENT', reports=False)

    response = client.post(WAVE_URL,
                           json={'assetids': [ready, offline, noagent, 999999]},
                           headers=auth_headers)
    assert response.status_code == 201, response.get_data(as_text=True)
    data = response.get_json()['data']
    assert data['queued'] == [ready]
    assert data['skipped'] == [
        {'assetid': offline, 'reason': 'offline'},
        {'assetid': noagent, 'reason': 'notreporting'},
        {'assetid': 999999, 'reason': 'not a computer'},
    ]


def test_wave_packages_once(client, db, app, auth_headers, monkeypatch):
    from plugins.endpoints.api import routes

    _publish_agent()
    ids = [_pc(db, f'PC-{n}') for n in range(3)]
    calls = []
    real = routes.prepare_agent_upgrade

    def counting():
        calls.append(1)
        return real()

    monkeypatch.setattr(routes, 'prepare_agent_upgrade', counting)
    response = client.post(WAVE_URL, json={'assetids': ids},
                           headers=auth_headers)
    assert response.status_code == 201
    assert response.get_json()['data']['queued'] == ids
    assert calls == [1]


def test_wave_progress_counts_enrolled_and_failed(client, db, app,
                                                  auth_headers):
    from plugins.endpoints.api.enrollment import issue_credential
    from plugins.endpoints.models import EndpointCommand

    _publish_agent()
    good = _pc(db, 'PC-GOOD')
    bad = _pc(db, 'PC-BAD')
    client.post(WAVE_URL, json={'assetids': [good, bad]}, headers=auth_headers)

    issue_credential(good, 'auto')
    EndpointCommand.query.filter_by(assetid=bad).one().status = \
        EndpointCommand.STATUS_FAILED
    db.session.commit()

    wave = _rows(client, auth_headers)['waves'][0]
    assert (wave['size'], wave['enrolled'], wave['failed'], wave['done']) == \
        (2, 1, 1, True)


def test_wave_without_ready_pcs_creates_nothing(client, db, app, auth_headers):
    from plugins.endpoints.models import EndpointMigrationWave

    offline = _pc(db, 'PC-OFFLINE', reportedago=timedelta(days=2))
    response = client.post(WAVE_URL, json={'assetids': [offline]},
                           headers=auth_headers)
    assert response.status_code == 200
    assert response.get_json()['data']['waveid'] is None
    assert EndpointMigrationWave.query.count() == 0


@pytest.mark.parametrize('body', [{}, {'next': 0}, {'next': 101},
                                  {'assetids': []}, {'assetids': ['1']},
                                  {'assetids': list(range(1, 102))}])
def test_wave_validation(client, db, app, auth_headers, body):
    assert client.post(WAVE_URL, json=body,
                       headers=auth_headers).status_code == 400