shopdb-flask/tests/test_plugins/test_endpoints_enroll_limits.py

"""The two bounds on POST /api/endpoints/enroll.

The route takes no auth on purpose - a PC installed without the shared key must
still be able to reach an admin - and it writes a row for every new
(hostname, serialnumber) pair. That is an unauthenticated caller writing to the
database in a loop, and it is also a denial of the enrollment workflow: junk
pending rows bury the real requests an admin has to approve.

Two bounds, and they answer different questions. The per-IP fixed-window limiter
bounds how FAST; the settings-backed cap bounds how MANY rows exist. Neither may
touch a PC already in the workflow, so the cap refuses only a hostname and
serial never seen before and an existing pending or approved row still answers
exactly as it did.
"""

import pytest

from shopdb.core.models import Asset, AssetType, Setting

from plugins.endpoints.api.enrollment import PENDING_CAP_SETTING
from plugins.endpoints.models import EndpointEnrollRequest


ENROLL_URL = '/api/endpoints/enroll'
NONCE = 'n' * 40


@pytest.fixture
def limits(app):
    """Both bounds on, at their shipped defaults, restored afterwards.

    The limiter counts in the cache extension, which is process-wide and
    outlives a test. Any other test in the same worker that posts to /enroll
    spends budget in the same window and IP bucket, so a count-exactly test has
    to start from a clean counter - and leave one behind, since it deliberately
    fills a bucket. Which tests share a worker depends on the xdist
    distribution, so without this the failure is intermittent and looks
    unrelated to whatever changed.
    """
    from shopdb.api import cache

    previous = {key: app.config.get(key) for key in (
        'ENDPOINTS_ENROLL_RATELIMIT_ENABLED',
        'ENDPOINTS_ENROLL_RATELIMIT_MAX',
        'ENDPOINTS_ENROLL_RATELIMIT_WINDOW_SECONDS')}
    app.config['ENDPOINTS_ENROLL_RATELIMIT_ENABLED'] = True
    with app.app_context():
        cache.clear()
    yield app
    with app.app_context():
        cache.clear()
    for key, value in previous.items():
        if value is None:
            app.config.pop(key, None)
        else:
            app.config[key] = value


def _computer(db, hostname, serialnumber):
    from plugins.computers.models import Computer

    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, assettypeid=atype.assettypeid,
                  serialnumber=serialnumber, isactive=True)
    db.session.add(asset)
    db.session.flush()
    db.session.add(Computer(assetid=asset.assetid, hostname=hostname))
    db.session.commit()
    return asset.assetid


def _body(hostname, serialnumber=None, nonce=NONCE):
    return {'hostname': hostname,
            'serialnumber': serialnumber or (hostname + '-SER'),
            'enrollnonce': nonce, 'agentversion': '1.20.0'}


def _cap(db, value):
    Setting.set(PENDING_CAP_SETTING, value, valuetype='integer',
                category='plugin.endpoints')
    db.session.commit()


def _pending(db, hostname, serialnumber=None, nonce=NONCE,
             status=EndpointEnrollRequest.STATUS_PENDING):
    from datetime import datetime, timezone

    now = datetime.now(timezone.utc).replace(tzinfo=None)
    row = EndpointEnrollRequest(
        hostname=hostname, serialnumber=serialnumber or (hostname + '-SER'),
        enrollnonce=nonce, status=status, requestedat=now)
    db.session.add(row)
    db.session.commit()
    return row


def _code(response):
    return response.get_json()['data']['error']['code']


# -- rate limit --------------------------------------------------------------

def test_the_limit_triggers_and_answers_the_standard_shape(client, db, limits):
    """Over budget is a 429 RATE_LIMITED, the same shape as every other
    limiter in the product."""
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_MAX'] = 3
    # An hour, so the fixed window cannot roll over mid-test and reset the
    # count. The test is about the budget, not about the clock.
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_WINDOW_SECONDS'] = 3600

    codes = []
    for index in range(5):
        response = client.post(ENROLL_URL, json=_body('BAY-%02d' % index))
        codes.append(response.status_code)

    assert codes[:3] == [202, 202, 202], codes
    assert codes[3:] == [429, 429], codes

    refused = client.post(ENROLL_URL, json=_body('BAY-99'))
    assert refused.status_code == 429
    assert _code(refused) == 'RATE_LIMITED'


def test_the_limit_is_off_when_the_config_says_so(client, db, limits):
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_ENABLED'] = False
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_MAX'] = 1

    for index in range(4):
        response = client.post(ENROLL_URL, json=_body('OFF-%02d' % index))
        assert response.status_code == 202, response.get_data(as_text=True)


# -- pending cap -------------------------------------------------------------

def test_the_cap_refuses_a_new_row(client, db, limits):
    """At the cap, a hostname and serial nobody has seen is told to try later,
    and no row is written for it."""
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_ENABLED'] = False
    _cap(db, 2)
    _pending(db, 'OLD-01')
    _pending(db, 'OLD-02')

    response = client.post(ENROLL_URL, json=_body('NEW-01'))

    assert response.status_code == 503, response.get_data(as_text=True)
    assert _code(response) == 'ENROLL_PENDING_FULL'
    assert EndpointEnrollRequest.query.filter_by(hostname='NEW-01').count() == 0
    assert EndpointEnrollRequest.query.count() == 2


def test_the_cap_still_serves_an_existing_pending_row(client, db, limits):
    """A PC already in the queue keeps asking and keeps being answered - the cap
    must not lock the fleet out of the workflow it is protecting."""
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_ENABLED'] = False
    _cap(db, 1)
    row = _pending(db, 'OLD-01')

    response = client.post(ENROLL_URL, json=_body('OLD-01'))

    assert response.status_code == 202, response.get_data(as_text=True)
    assert response.get_json()['data']['status'] == 'pending'
    db.session.refresh(row)
    assert row.lastseenat is not None


def test_the_cap_still_issues_to_an_approved_row(client, db, limits):
    """An approved request hands out its credential at the cap, because that
    row is the admin's decision and refusing it would strand the PC."""
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_ENABLED'] = False
    assetid = _computer(db, 'BAY-07', 'BAY-07-SER')
    _cap(db, 1)
    _pending(db, 'FILLER-01')
    approved = _pending(db, 'BAY-07',
                        status=EndpointEnrollRequest.STATUS_APPROVED)
    approved.matchedassetid = assetid
    db.session.commit()

    response = client.post(ENROLL_URL, json=_body('BAY-07'))

    assert response.status_code == 201, response.get_data(as_text=True)
    assert response.get_json()['data']['credential'].startswith('sdbe_')


def test_the_cap_setting_is_honoured(client, db, limits):
    """Raise the setting and the same request that was refused goes through;
    0 removes the cap entirely."""
    limits.config['ENDPOINTS_ENROLL_RATELIMIT_ENABLED'] = False
    _cap(db, 1)
    _pending(db, 'OLD-01')

    refused = client.post(ENROLL_URL, json=_body('NEW-01'))
    assert refused.status_code == 503

    _cap(db, 5)
    allowed = client.post(ENROLL_URL, json=_body('NEW-01'))
    assert allowed.status_code == 202, allowed.get_data(as_text=True)

    _cap(db, 0)
    uncapped = client.post(ENROLL_URL, json=_body('NEW-02'))
    assert uncapped.status_code == 202, uncapped.get_data(as_text=True)


def test_the_cap_is_editable_through_the_enrollment_settings_route(
        client, db, auth_headers):
    """A setting nobody can change is a constant with extra steps."""
    response = client.put('/api/endpoints/enrollment/settings',
                          json={'maxpendingenrollments': 42},
                          headers=auth_headers)

    assert response.status_code == 200, response.get_data(as_text=True)
    assert response.get_json()['data']['maxpendingenrollments'] == 42
    assert int(Setting.get(PENDING_CAP_SETTING)) == 42

    bad = client.put('/api/endpoints/enrollment/settings',
                     json={'maxpendingenrollments': -1}, headers=auth_headers)
    assert bad.status_code == 400