"""GE-Enforce payload-download hardening: rate limit + served-size ceiling.

GET /api/geenforce/payload/<sha256> is reachable with only a read-only
geenforce.fetch token, so a leaked display token must not be able to hammer it
or pull unbounded bytes. These tests lock the two bounds added to the route:
a per-IP fixed-window rate limit (same cache extension as the login limiter)
and a size cap that refuses (413) a blob/inline payload over the configured max.
"""

import hashlib
import io
from contextlib import contextmanager

from plugins.geenforce import service
from plugins.geenforce.models import ManifestBlob


def _fetch_key(client, auth_headers):
    resp = client.post('/api/apitokens',
                       json={'name': 'svc', 'scopes': ['geenforce.fetch']},
                       headers=auth_headers)
    assert resp.status_code == 201, resp.get_json()
    return {'X-API-Key': resp.get_json()['data']['secret']}


def _store(app, raw, filename='setup.exe'):
    with app.app_context():
        sha = service.store_blob(raw, filename, 'application/octet-stream')
        service.db.session.commit()
    return sha


@contextmanager
def _config(app, **overrides):
    """Temporarily override app.config keys (the app fixture is session-scoped,
    so every key is restored to its prior value/absence on exit)."""
    missing = object()
    saved = {key: app.config.get(key, missing) for key in overrides}
    app.config.update(overrides)
    try:
        yield
    finally:
        for key, value in saved.items():
            if value is missing:
                app.config.pop(key, None)
            else:
                app.config[key] = value


def test_blob_within_cap_is_served(client, app, auth_headers):
    raw = b'x' * 64
    sha = _store(app, raw, 'ok.exe')
    headers = _fetch_key(client, auth_headers)
    with _config(app, GEENFORCE_PAYLOAD_MAX_BYTES=1024):
        resp = client.get(f'/api/geenforce/payload/{sha}', headers=headers)
    assert resp.status_code == 200
    assert resp.data == raw


def test_oversized_blob_refused_413(client, app, auth_headers):
    raw = b'x' * 64
    sha = _store(app, raw, 'big.exe')
    headers = _fetch_key(client, auth_headers)
    with _config(app, GEENFORCE_PAYLOAD_MAX_BYTES=16):
        resp = client.get(f'/api/geenforce/payload/{sha}', headers=headers)
    assert resp.status_code == 413
    # The blob registry row still exists; the cap gates delivery, not storage.
    with app.app_context():
        assert service.db.session.get(ManifestBlob, sha) is not None


def test_oversized_inline_payload_refused_413(client, app, db, auth_headers):
    scopeid = client.post(
        '/api/geenforce/scopes',
        json={'scopename': 'gea-shopfloor-cmm', 'phase': 'runtime'},
        headers=auth_headers).get_json()['data']['scopeid']
    entryid = client.post(
        f'/api/geenforce/scopes/{scopeid}/entries',
        json={'Name': 'eDNC config', 'Type': 'PS1', 'Script': 's.ps1'},
        headers=auth_headers).get_json()['data']['entryid']
    content = b'inline config bytes here'
    up = client.post(f'/api/geenforce/entries/{entryid}/payload',
                     data={'file': (io.BytesIO(content), 'config.reg')},
                     content_type='multipart/form-data', headers=auth_headers)
    assert up.status_code == 201, up.get_json()
    sha = hashlib.sha256(content).hexdigest()

    headers = _fetch_key(client, auth_headers)
    with _config(app, GEENFORCE_PAYLOAD_MAX_BYTES=4):
        resp = client.get(f'/api/geenforce/payload/{sha}', headers=headers)
    assert resp.status_code == 413


def test_payload_download_is_rate_limited(client, app, auth_headers):
    sha = _store(app, b'small blob')
    key = _fetch_key(client, auth_headers)
    # A unique caller IP isolates this bucket from other tests' shared counter.
    headers = {**key, 'X-Forwarded-For': '203.0.113.201'}
    with _config(app, GEENFORCE_PAYLOAD_RATELIMIT_MAX=3,
                 GEENFORCE_PAYLOAD_RATELIMIT_WINDOW_SECONDS=300):
        for _ in range(3):
            ok = client.get(f'/api/geenforce/payload/{sha}', headers=headers)
            assert ok.status_code == 200, ok.get_json()
        # The 4th request in the same window is over budget.
        over = client.get(f'/api/geenforce/payload/{sha}', headers=headers)
    assert over.status_code == 429


def test_rate_limit_is_per_ip(client, app, auth_headers):
    """A second caller (different X-Forwarded-For) has its own budget: one IP
    being throttled must not throttle everyone."""
    sha = _store(app, b'per-ip blob')
    key = _fetch_key(client, auth_headers)
    hot = {**key, 'X-Forwarded-For': '203.0.113.202'}
    cool = {**key, 'X-Forwarded-For': '203.0.113.203'}
    with _config(app, GEENFORCE_PAYLOAD_RATELIMIT_MAX=1,
                 GEENFORCE_PAYLOAD_RATELIMIT_WINDOW_SECONDS=300):
        assert client.get(f'/api/geenforce/payload/{sha}',
                          headers=hot).status_code == 200
        assert client.get(f'/api/geenforce/payload/{sha}',
                          headers=hot).status_code == 429
        # Fresh IP still gets its first request through.
        assert client.get(f'/api/geenforce/payload/{sha}',
                          headers=cool).status_code == 200


def test_rate_limit_can_be_disabled(client, app, auth_headers):
    sha = _store(app, b'unthrottled blob')
    key = _fetch_key(client, auth_headers)
    headers = {**key, 'X-Forwarded-For': '203.0.113.204'}
    with _config(app, GEENFORCE_PAYLOAD_RATELIMIT_ENABLED=False,
                 GEENFORCE_PAYLOAD_RATELIMIT_MAX=1):
        for _ in range(5):
            assert client.get(f'/api/geenforce/payload/{sha}',
                              headers=headers).status_code == 200
