truf-server/tests/edge_e2e_client.py

"""Same-IP HTTPS assertions for the standalone edge E2E orchestrator."""

import sys

sys.dont_write_bytecode = True

import base64
import hashlib
import http.client
import json
import os
from pathlib import Path
import re
import runpy
import socket
import ssl
import time
import uuid
from urllib.parse import urlencode


APP = Path('/opt/truf/app')
TESTS = Path('/opt/truf/tests')
CERTIFICATE = Path('/opt/truf/tests/fixtures/worker_tls_cert.pem')
MAX_BODY = 1024 * 1024
FAILURE_STAGE = 'startup'
SECURITY_HEADERS = {
    'cache-control': 'no-store',
    'referrer-policy': 'same-origin',
    'content-security-policy': "default-src 'none'",
    'x-content-type-options': 'nosniff',
    'strict-transport-security': 'max-age=',
}


def require(condition, label):
    if not condition:
        raise RuntimeError('standalone edge E2E client: ' + label)


def environment(name, minimum=1):
    value = str(os.environ.get(name) or '')
    require(len(value) >= minimum, 'missing test environment')
    return value


PREFIX = environment('TRUF_ADMIN_PREFIX', 64)
ADMIN_USER = environment('TRUF_ADMIN_USER')
ADMIN_PASSWORD = environment('TRUF_EDGE_E2E_ADMIN_PASSWORD', 16)
CONNECT_HOST = environment('TRUF_EDGE_E2E_CONNECT_HOST')
DB_URL = f'postgresql://truf@{CONNECT_HOST}:55433/edge_e2e'
TOKENS = {
    name: environment(variable, 32)
    for name, variable in {
        'good': 'TRUF_EDGE_E2E_GOOD_TOKEN',
        'wrong': 'TRUF_EDGE_E2E_WRONG_TOKEN',
        'revoked': 'TRUF_EDGE_E2E_REVOKED_TOKEN',
    }.items()
}


def basic(user=ADMIN_USER, password=ADMIN_PASSWORD):
    payload = base64.b64encode(f'{user}:{password}'.encode('utf-8')).decode('ascii')
    return 'Basic ' + payload


def request(method, path, *, headers=None, body=None, expected=None, secure=True):
    context = ssl.create_default_context(cafile=str(CERTIFICATE))
    connection = http.client.HTTPSConnection(
        'localhost', 443, context=context, timeout=10,
    )
    connection._create_connection = lambda _address, timeout=None, source_address=None: (
        socket.create_connection((CONNECT_HOST, 443), timeout, source_address)
    )
    try:
        connection.request(method, path, body=body, headers=headers or {})
        response = connection.getresponse()
        payload = response.read(MAX_BODY + 1)
        require(len(payload) <= MAX_BODY, 'response body bound')
        response_headers = {name.lower(): value for name, value in response.getheaders()}
    finally:
        connection.close()
    if expected is not None:
        require(response.status == expected, f'{method} {path} status {response.status}')
    if secure:
        for name, fragment in SECURITY_HEADERS.items():
            require(fragment in response_headers.get(name, ''), f'{name} on {path}')
    return response.status, response_headers, payload


def json_request(path, token, request_id):
    from edge_e2e_backend import worker_build

    payload = json.dumps(
        {'request_id': request_id, 'build': worker_build()},
        ensure_ascii=True, sort_keys=True, separators=(',', ':'),
    ).encode('ascii')
    return request(
        'POST', path,
        headers={
            'Authorization': 'Bearer ' + token,
            'Content-Type': 'application/json',
            'X-Truf-Admin-Edge': 'spoofed-edge-marker',
            'X-Truf-Admin-Operator': 'spoofed-operator',
        },
        body=payload,
        secure=False,
    )


def db_rows(statement, values=()):
    from scanner_db import ScannerDB

    db = ScannerDB(db_url=DB_URL, initialize=False)
    require(db.enabled and db.conn.is_postgres, 'client PostgreSQL connection')
    try:
        rows = db.conn.execute(statement, values).fetchall()
        db.conn.commit()
        return [tuple(dict(row).values()) for row in rows]
    finally:
        db.close()


def assignment_state(reservation_id):
    return {
        'reservation': db_rows(
            '''SELECT id, queue_id, state, assignment_kind, remote_user_id,
                      remote_device_id, remote_resolution_kind, remote_receipt_id,
                      remote_expires_at, updated_at
               FROM result_reservations WHERE id = ?''',
            (reservation_id,),
        ),
        'queue': db_rows(
            '''SELECT q.id, q.status, q.attempts, q.lease_token,
                      q.current_result_reservation_id, q.claim_event_id, q.updated_at
               FROM target_queue q JOIN result_reservations r ON r.queue_id = q.id
               WHERE r.id = ?''',
            (reservation_id,),
        ),
        'capacity': db_rows(
            '''SELECT bundle_items, bundle_bytes, projection_items, projection_bytes,
                      keycheck_items, keycheck_bytes, updated_at
               FROM pipeline_capacity WHERE id = 1''',
        ),
    }


def reservation_id():
    rows = db_rows(
        '''SELECT r.id FROM result_reservations r
           JOIN remote_worker_devices d ON d.id = r.remote_device_id
           WHERE d.device_key = 'edge-device-good'
           ORDER BY r.id DESC LIMIT 1''',
    )
    require(len(rows) == 1 and int(rows[0][0]) > 0, 'good assignment lookup')
    return int(rows[0][0])


def worker_status(reservation, token=TOKENS['good'], expected=200):
    return request(
        'GET', f'/api/v1/worker/assignments/{reservation}',
        headers={
            'Authorization': 'Bearer ' + token,
            'X-Truf-Admin-Edge': 'spoofed-edge-marker',
            'X-Truf-Admin-Operator': 'spoofed-operator',
        }, expected=expected, secure=False,
    )


def admin_request(
    path='', *, authorization=True, expected=200, method='GET', body=None,
    origin=None, secure=True,
):
    headers = {'X-Truf-Admin-Operator': 'spoofed-operator'}
    if authorization:
        headers['Authorization'] = basic()
    if origin is not None:
        headers['Origin'] = origin
    if body is not None:
        headers['Content-Type'] = 'application/x-www-form-urlencoded'
    return request(
        method, f'/{PREFIX}/{path}', headers=headers, body=body, expected=expected,
        secure=secure,
    )


def baseline():
    global FAILURE_STAGE
    FAILURE_STAGE = 'public_routes'
    checked = 0
    for path in (
        '/unknown', '/dashboard', '/metrics', '/admin-internal',
        '/api/v1/private', '/api/v1/worker/private',
    ):
        request('GET', path, expected=404, secure=False)
        checked += 1

    FAILURE_STAGE = 'worker_challenge'
    missing_payload = b'{"build":{},"request_id":"00000000000000000000000000000000"}'
    status, headers, _ = request(
        'POST', '/api/v1/worker/claim',
        headers={'Content-Type': 'application/json'}, body=missing_payload, expected=401,
        secure=False,
    )
    require(headers.get('www-authenticate') == 'Bearer', 'missing worker Bearer challenge')
    checked += 1

    FAILURE_STAGE = 'admin_challenge'
    for path in ('', 'admin.css'):
        status, headers, _ = admin_request(
            path, authorization=False, expected=401, secure=False,
        )
        require(
            headers.get('www-authenticate') == 'Basic realm="truf-admin"',
            'missing admin Basic challenge',
        )
        checked += 1

    FAILURE_STAGE = 'worker_auth'
    status, _, _ = json_request(
        '/api/v1/worker/claim', TOKENS['revoked'], '1' * 32,
    )
    require(status == 401, 'revoked worker token')
    checked += 1

    FAILURE_STAGE = 'worker_api'
    status, _, payload = json_request(
        '/api/v1/worker/claim', TOKENS['good'], '2' * 32,
    )
    require(status == 201, 'valid worker claim')
    value = json.loads(payload.decode('utf-8'))
    reservation = int(value['assignment']['reservation']['reservation_id'])
    require(reservation > 0, 'valid worker reservation')
    checked += 1

    before = assignment_state(reservation)
    worker_status(reservation, TOKENS['wrong'], expected=404)
    after = assignment_state(reservation)
    require(after == before, 'wrong-device request changed authoritative assignment state')
    checked += 1
    worker_status(reservation)
    checked += 1

    FAILURE_STAGE = 'admin_page_render'
    _, _, page = admin_request()
    checked += 1
    require(b'Workers / Dispatch' in page, 'protected admin page')
    match = re.search(rb'name="csrf_token" value="([A-Za-z0-9_-]{32,128})"', page)
    require(match is not None, 'admin CSRF token')
    csrf = match.group(1).decode('ascii')
    FAILURE_STAGE = 'admin_asset'
    _, _, css = admin_request('admin.css')
    require(b'color-scheme' in css, 'protected admin asset')
    checked += 1

    users_before = db_rows('SELECT user_key, active_assignment_cap, disabled_at FROM remote_worker_users ORDER BY id')
    FAILURE_STAGE = 'admin_cross_site'
    cross_site_operation_id = str(uuid.uuid4())
    cross_site = urlencode({
        'csrf_token': csrf, 'user_key': 'cross-site-user',
        'active_assignment_cap': '2', 'operation_id': cross_site_operation_id,
    }).encode('ascii')
    admin_request(
        'users/create', method='POST', body=cross_site,
        origin='https://cross-site.invalid', expected=403,
    )
    require(
        db_rows('SELECT user_key, active_assignment_cap, disabled_at FROM remote_worker_users ORDER BY id')
        == users_before,
        'cross-site mutation changed authoritative state',
    )
    checked += 1

    FAILURE_STAGE = 'admin_same_site'
    same_site_operation_id = str(uuid.uuid4())
    same_site = urlencode({
        'csrf_token': csrf, 'user_key': 'same-site-user',
        'active_assignment_cap': '2', 'operation_id': same_site_operation_id,
    }).encode('ascii')
    _, mutation_headers, _ = admin_request(
        'users/create', method='POST', body=same_site,
        origin='https://localhost', expected=303,
    )
    require(mutation_headers.get('location') == '../', 'same-origin mutation redirect')
    require(
        db_rows("SELECT user_key, active_assignment_cap FROM remote_worker_users WHERE user_key = 'same-site-user'")
        == [('same-site-user', 2)],
        'same-origin mutation persistence',
    )
    FAILURE_STAGE = 'admin_operation_actor'
    require(
        db_rows(
            '''SELECT actor, action, target_kind, target_ref, status
               FROM runtime_operations WHERE operation_id = ?''',
            (same_site_operation_id,),
        ) == [(ADMIN_USER, 'workers.user.create', 'worker-admin', 'same-site-user', 'succeeded')],
        'same-origin durable operation actor attribution',
    )
    FAILURE_STAGE = 'admin_audit_actor'
    require(
        db_rows(
            '''SELECT actor, action, target_kind, target_ref, result
               FROM runtime_audit_events WHERE operation_id = ? ORDER BY id''',
            (same_site_operation_id,),
        ) == [
            (ADMIN_USER, 'workers.user.create', 'worker-admin', 'same-site-user', 'accepted'),
            (ADMIN_USER, 'workers.user.create', 'worker-admin', 'same-site-user', 'succeeded'),
        ],
        'same-origin durable audit actor attribution',
    )
    checked += 1

    FAILURE_STAGE = 'admin_route_crawl'
    admin_routes = (
        'overview', 'search', 'supervisor', 'logs', 'config', 'secrets',
        'files', 'operations', 'audit',
        f'operations/{same_site_operation_id}',
    )
    for route in admin_routes:
        route_label = (
            'operation_detail' if route.startswith('operations/') else route
        )
        FAILURE_STAGE = 'admin_route_' + route_label
        route_status, _, route_page = admin_request(route, expected=None)
        FAILURE_STAGE = f'admin_route_{route_label}_{route_status}'
        require(route_status in (200, 503), 'protected admin route status')
        if route_status == 200:
            FAILURE_STAGE = 'admin_route_' + route_label + '_body'
            require(b'<!doctype html>' in route_page.lower(), 'protected admin route body')
        checked += 1

    FAILURE_STAGE = 'admin_worker_status'
    worker_status(reservation)
    checked += 1
    return {
        'mode': 'baseline', 'checked_responses': checked,
        'reservation_id': reservation, 'operation_id': same_site_operation_id,
    }


def bad_auth():
    attempts = (
        ('truf-admin', 'definitely-wrong-one', '198.51.100.17'),
        ('not-the-admin', 'definitely-wrong-two', '2001:db8::17'),
    )
    for user, password, spoofed in attempts:
        request(
            'GET', f'/{PREFIX}/',
            headers={
                'Authorization': basic(user, password),
                'X-Forwarded-For': spoofed,
                'Forwarded': f'for="[{spoofed}]";proto=http;host=spoofed.invalid',
            },
            expected=401,
        )
    return {'mode': 'bad-auth', 'attempts': len(attempts)}


def assert_ban():
    global FAILURE_STAGE
    FAILURE_STAGE = 'ban_lookup'
    reservation = reservation_id()
    FAILURE_STAGE = 'ban_admin'
    admin_request(expected=403)
    FAILURE_STAGE = 'ban_worker'
    worker_status(reservation)
    return {'mode': 'assert-ban', 'reservation_id': reservation}


def assert_unban():
    global FAILURE_STAGE
    FAILURE_STAGE = 'unban_lookup'
    reservation = reservation_id()
    FAILURE_STAGE = 'unban_admin'
    _, _, page = admin_request()
    require(b'Workers / Dispatch' in page, 'admin restoration')
    FAILURE_STAGE = 'unban_worker'
    worker_status(reservation)
    return {'mode': 'assert-unban', 'reservation_id': reservation}


def probe():
    status, headers, _ = request('GET', '/edge-e2e-probe', secure=False)
    missing = [
        name for name, fragment in SECURITY_HEADERS.items()
        if fragment not in headers.get(name, '')
    ]
    return {
        'mode': 'probe',
        'ready': status == 404 and not missing,
        'status': status,
        'missing': missing,
    }


def main():
    require(
        sys.platform == 'linux' and sys.flags.isolated and sys.flags.no_site
        and sys.flags.dont_write_bytecode,
        'isolated Linux Python required',
    )
    sys.path.insert(0, str(APP))
    sys.path.insert(0, str(TESTS))
    bootstrap = runpy.run_path(str(APP / 'child_bootstrap.py'))
    bootstrap['_enable_dependency_paths']('supervisor')
    modes = {
        'baseline': baseline, 'bad-auth': bad_auth, 'assert-ban': assert_ban,
        'assert-unban': assert_unban, 'probe': probe,
    }
    require(len(sys.argv) == 2 and sys.argv[1] in modes, 'mode')
    try:
        result = modes[sys.argv[1]]()
    except Exception:
        stage = FAILURE_STAGE
        if not re.fullmatch(r'[a-z][a-z0-9_]{0,39}', stage):
            stage = 'runtime'
        print(json.dumps({
            'mode': sys.argv[1], 'failure_stage': stage,
            'tls': 'validated-localhost-certificate',
        }, ensure_ascii=True, sort_keys=True, separators=(',', ':')))
        return 1
    result['tls'] = 'validated-localhost-certificate'
    print(json.dumps(result, ensure_ascii=True, sort_keys=True, separators=(',', ':')))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())