rapp-playwright/client/rapp_playwright.py

#!/usr/bin/env python3
"""Dependency-free client for the rapp-playwright v1 job API (Python 3.10+)."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class ClientError(Exception):
    pass


class Client:
    def __init__(self, url):
        parsed = urllib.parse.urlsplit(url)
        if parsed.scheme not in ('http', 'https') or parsed.username or parsed.password or parsed.query or parsed.fragment:
            raise ClientError('invalid service URL')
        if parsed.scheme == 'http' and parsed.hostname not in ('localhost', '127.0.0.1', '::1'):
            raise ClientError('HTTP requires a loopback tunnel; remote endpoints require HTTPS')
        self.url = url.rstrip('/')
        self.opener = urllib.request.build_opener(NoRedirect)

    def request(self, path, method='GET', body=None, token=None, binary=False):
        headers = {'X-Rapp-Protocol': '1'}
        if token:
            headers['Authorization'] = 'Bearer ' + token
        data = None
        if body is not None:
            headers['Content-Type'] = 'application/json'
            data = json.dumps(body).encode()
        req = urllib.request.Request(self.url + path, data=data, headers=headers, method=method)
        try:
            with self.opener.open(req, timeout=30) as response:
                data = response.read(21 * 1024 * 1024)
                if len(data) >= 21 * 1024 * 1024:
                    raise ClientError('response exceeds artifact limit')
                return data if binary else json.loads(data)
        except urllib.error.HTTPError as exc:
            try:
                detail = json.loads(exc.read(32768))
            except (ValueError, OSError):
                detail = {'error': 'http_error'}
            raise ClientError(f'HTTP {exc.code}: {detail.get("error", "service_unavailable")}') from None
        except (urllib.error.URLError, TimeoutError, OSError):
            raise ClientError('service_unreachable: establish the documented tunnel and retry') from None


def save(path, data):
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, 'w') as stream:
        json.dump(data, stream, indent=2)
        stream.write('\n')


def handle(path):
    data = json.loads(path.read_text())
    if not re.fullmatch(r'[a-f0-9-]{36}', data['id']) or not re.fullmatch(r'[a-f0-9]{64}', data['token']):
        raise ClientError('invalid local run handle')
    return data


def collect(client, run, out):
    receipt = client.request('/v1/runs/' + run['id'], token=run['token'])
    downloads = []
    for artifact in receipt.get('artifacts', []):
        name = artifact['name']
        if not re.fullmatch(r'[a-z0-9-]+\.png', name):
            raise ClientError('invalid artifact name')
        data = client.request('/v1/runs/' + run['id'] + '/artifacts/' + name, token=run['token'], binary=True)
        if len(data) != artifact['bytes']:
            raise ClientError('artifact size mismatch')
        fd = os.open(out / name, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, 'wb') as stream:
            stream.write(data)
        downloads.append({'name': name, 'bytes': len(data), 'sha256': hashlib.sha256(data).hexdigest()})
    receipt['downloads'] = downloads
    save(out / 'receipt.json', receipt)
    return receipt


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--url', default=os.environ.get('RAPP_PLAYWRIGHT_URL', 'http://127.0.0.1:18081'))
    subs = parser.add_subparsers(dest='command', required=True)
    subs.add_parser('preflight')
    for cmd in ('trial', 'run'):
        sub = subs.add_parser(cmd)
        sub.add_argument('--out', type=Path, required=True, help='New private artifact directory')
        if cmd == 'run':
            sub.add_argument('--spec', type=Path, required=True)
    for cmd in ('status', 'cancel', 'download'):
        sub = subs.add_parser(cmd)
        sub.add_argument('--handle', type=Path, required=True)
    args = parser.parse_args()
    client = Client(args.url)
    if args.command == 'preflight':
        print(json.dumps(client.request('/v1/preflight'), indent=2))
        return 0
    if args.command in ('status', 'cancel', 'download'):
        run = handle(args.handle)
        if args.command == 'download':
            result = collect(client, run, args.handle.parent)
        else:
            result = client.request('/v1/runs/' + run['id'], method='DELETE' if args.command == 'cancel' else 'GET', token=run['token'])
        print(json.dumps(result, indent=2))
        return 0
    preflight = client.request('/v1/preflight')
    if not preflight['available'] or preflight['api_version'] != '1':
        raise ClientError('service_not_ready_or_incompatible')
    spec = client.request('/v1/trial') if args.command == 'trial' else json.loads(args.spec.read_text())
    args.out.mkdir(parents=True, mode=0o700, exist_ok=False)
    save(args.out / 'preflight.json', preflight)
    save(args.out / 'invocation.json', {'argv': sys.argv, 'service_url': args.url})
    run = client.request('/v1/runs', 'POST', spec)
    save(args.out / 'run.json', run)  # Contains the capability; never print or publish.
    deadline = time.monotonic() + spec.get('timeout_ms', 60000) / 1000 + 520
    try:
        while True:
            result = client.request('/v1/runs/' + run['id'], token=run['token'])
            if result['status'] in ('passed', 'failed', 'cancelled', 'interrupted'):
                break
            if time.monotonic() > deadline:
                raise ClientError('client_wait_timeout: run cancellation requested')
            time.sleep(0.25)
    except (KeyboardInterrupt, ClientError):
        try:
            client.request('/v1/runs/' + run['id'], 'DELETE', token=run['token'])
        except ClientError:
            pass  # Server time budget still bounds disconnected clients.
        raise
    result = collect(client, run, args.out)
    print(json.dumps({'id': result['id'], 'status': result['status'], 'error': result['error'],
                      'receipt': str(args.out / 'receipt.json'), 'artifacts': result['downloads']}, indent=2))
    return 0 if result['status'] == 'passed' else 2 if result['error'] == 'ui_assertion_failed' else 3


if __name__ == '__main__':
    try:
        sys.exit(main())
    except (ClientError, ValueError, OSError, KeyError) as exc:
        print(json.dumps({'error': str(exc)}), file=sys.stderr)
        sys.exit(3)
    except KeyboardInterrupt:
        print(json.dumps({'error': 'cancelled_by_client'}), file=sys.stderr)
        sys.exit(4)