"""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


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

    _publish_agent()
    assetid = _pc(db, 'PC-ROLLBACK')
    issue_credential(assetid, 'auto')
    db.session.commit()
    response = client.post(
        f'/api/endpoints/enrollment/assets/{assetid}/revoke',
        json={'blockreenroll': True, 'reason': 'rollback'},
        headers=auth_headers)
    assert response.status_code == 200, response.get_data(as_text=True)
    # Its agent keeps asking; that request must not turn the row into pending.
    db.session.add(EndpointEnrollRequest(
        hostname='PC-ROLLBACK', serialnumber='SER', enrollnonce='n' * 40,
        status='pending', reason='reenroll blocked', matchedassetid=assetid,
        requestedat=_now(), lastseenat=_now()))
    db.session.commit()

    row = _by_host(_rows(client, auth_headers))['PC-ROLLBACK']
    assert row['state'] == 'blocked'
    assert row['blockers'] == ['reenroll blocked']

    wave = client.post(WAVE_URL, json={'assetids': [assetid]},
                       headers=auth_headers).get_json()['data']
    assert wave['skipped'] == [{'assetid': assetid,
                                'reason': 'reenroll blocked'}]

    client.post(f'/api/endpoints/enrollment/assets/{assetid}/allow-reenroll',
                headers=auth_headers)
    assert _by_host(_rows(client, auth_headers))['PC-ROLLBACK']['state'] \
        == 'pending'


@pytest.mark.parametrize('message,detail', [
    ('Script download failed: The SSL connection could not be established, '
     'see inner exception.', None),
    ('Script download failed', {'error': 'The remote certificate is invalid '
                                'because of errors in the certificate chain: '
                                'UntrustedRoot'}),
    ('Invoke-WebRequest: Could not establish trust relationship for the '
     'SSL/TLS secure channel.', None),
])
def test_a_tls_trust_failure_blocks_the_pc(client, db, app, auth_headers,
                                           message, detail):
    import json

    from plugins.endpoints.models import EndpointCommand

    _publish_agent()
    assetid = _pc(db, 'PC-UNTRUSTED')
    otherid = _pc(db, 'PC-OTHERFAIL')
    agentregistry.register_current_agent_package()
    scriptid = agentupgradescript.seed_upgrade_script()['scriptid']
    db.session.add(EndpointCommand(
        assetid=assetid, action='run_script', target=f'script:{scriptid}',
        status=EndpointCommand.STATUS_FAILED, resultmessage=message,
        resultjson=json.dumps(detail) if detail is not None else None,
        completedat=_now()))
    db.session.add(EndpointCommand(
        assetid=otherid, action='run_script', target=f'script:{scriptid}',
        status=EndpointCommand.STATUS_FAILED,
        resultmessage='Script exited with code 1', completedat=_now()))
    db.session.commit()

    rows = _by_host(_rows(client, auth_headers))
    assert rows['PC-UNTRUSTED']['state'] == 'blocked'
    assert rows['PC-UNTRUSTED']['blockers'] == [
        'server certificate not trusted by this PC']
    assert rows['PC-OTHERFAIL']['state'] == 'ready'


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

    _publish_agent()
    assetid = _pc(db, 'PC-FIXED')
    agentregistry.register_current_agent_package()
    scriptid = agentupgradescript.seed_upgrade_script()['scriptid']
    db.session.add(EndpointCommand(
        assetid=assetid, action='run_script', target=f'script:{scriptid}',
        status=EndpointCommand.STATUS_FAILED,
        resultmessage='RemoteCertificateChainErrors', completedat=_now()))
    db.session.flush()
    db.session.add(EndpointCommand(
        assetid=assetid, action='run_script', target='script:999',
        status=EndpointCommand.STATUS_SUCCEEDED, completedat=_now()))
    db.session.commit()

    assert _by_host(_rows(client, auth_headers))['PC-FIXED']['state'] == 'ready'


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

    _publish_agent()
    trustonly = _pc(db, 'PC-TRUST')
    alsooffline = _pc(db, 'PC-TRUST-OFFLINE', reportedago=timedelta(days=2))
    agentregistry.register_current_agent_package()
    scriptid = agentupgradescript.seed_upgrade_script()['scriptid']
    for assetid in (trustonly, alsooffline):
        db.session.add(EndpointCommand(
            assetid=assetid, action='run_script', target=f'script:{scriptid}',
            status=EndpointCommand.STATUS_FAILED,
            resultmessage='UntrustedRoot', completedat=_now()))
    db.session.commit()

    rows = _by_host(_rows(client, auth_headers))
    assert rows['PC-TRUST']['blockers'] == [
        'server certificate not trusted by this PC']
    assert rows['PC-TRUST-OFFLINE']['blockers'] == [
        'offline', 'server certificate not trusted by this PC']

    nextwave = client.post(WAVE_URL, json={'next': 10}, headers=auth_headers)
    assert nextwave.status_code == 200, nextwave.get_data(as_text=True)
    assert nextwave.get_json()['data']['queued'] == []
    assert nextwave.get_json()['data']['retried'] == []
    assert EndpointMigrationWave.query.count() == 0

    response = client.post(WAVE_URL,
                           json={'assetids': [trustonly, alsooffline]},
                           headers=auth_headers)
    assert response.status_code == 201, response.get_data(as_text=True)
    data = response.get_json()['data']
    assert data['queued'] == [trustonly]
    assert data['retried'] == [trustonly]
    assert data['skipped'] == [{
        'assetid': alsooffline,
        'reason': 'offline, server certificate not trusted by this PC'}]
    assert EndpointMigrationWave.query.one().assetidlist == [trustonly]
    assert _by_host(_rows(client, auth_headers))['PC-TRUST']['state'] \
        == 'upgrading'


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

    _publish_agent()
    assetid = _pc(db, 'PC-TRUST')
    agentregistry.register_current_agent_package()
    scriptid = agentupgradescript.seed_upgrade_script()['scriptid']
    db.session.add(EndpointCommand(
        assetid=assetid, action='run_script', target=f'script:{scriptid}',
        status=EndpointCommand.STATUS_FAILED,
        resultmessage='UntrustedRoot', completedat=_now()))
    db.session.commit()
    monkeypatch.setattr(enrollment, 'ALLOWED_ACTIONS_FOR',
                        lambda assetid=None: set())

    data = client.post(WAVE_URL, json={'assetids': [assetid]},
                       headers=auth_headers).get_json()['data']
    assert data['queued'] == []
    assert data['skipped'] == [{
        'assetid': assetid,
        'reason': 'server certificate not trusted by this PC, '
                  'run_script not permitted by policy'}]


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

    _publish_agent()
    upgraded = _pc(db, 'PC-UPGRADED')
    retried = _pc(db, 'PC-RETRIED')
    stuck = _pc(db, 'PC-STUCK')
    first = client.post(WAVE_URL, json={'assetids': [upgraded, retried, stuck]},
                        headers=auth_headers)
    assert first.status_code == 201, first.get_data(as_text=True)
    for command in EndpointCommand.query.all():
        command.status = EndpointCommand.STATUS_FAILED
    db.session.commit()
    assert _rows(client, auth_headers)['waves'][0]['failed'] == 3

    second = client.post(WAVE_URL, json={'assetids': [upgraded, retried]},
                         headers=auth_headers)
    assert second.status_code == 201, second.get_data(as_text=True)
    # The later wave replaced the agent on PC-UPGRADED, which restarted before
    # reporting, so its command expired; it has not enrolled yet. PC-RETRIED's
    # push reported success and is waiting to enroll on its old version string.
    latest = {command.assetid: command for command in EndpointCommand.query
              .filter_by(status=EndpointCommand.STATUS_PENDING).all()}
    latest[upgraded].status = EndpointCommand.STATUS_EXPIRED
    latest[retried].status = EndpointCommand.STATUS_SUCCEEDED
    EndpointState.query.filter_by(assetid=upgraded).one().agentversion = '1.20.0'
    db.session.commit()

    failed = {wave['waveid']: wave['failed']
              for wave in _rows(client, auth_headers)['waves']}
    assert failed == {first.get_json()['data']['waveid']: 1,
                      second.get_json()['data']['waveid']: 0}

    # A newer failure counts again, in every wave the PC belongs to.
    latest[retried].status = EndpointCommand.STATUS_FAILED
    db.session.commit()
    failed = {wave['waveid']: wave['failed']
              for wave in _rows(client, auth_headers)['waves']}
    assert failed == {first.get_json()['data']['waveid']: 2,
                      second.get_json()['data']['waveid']: 1}