pigway-console-adapter-raspberry-pi/plugin/worker.py

#!/usr/bin/env python3
"""Official Raspberry Pi adapter. This process alone knows Agent endpoints/authentication."""
import concurrent.futures
import json
import math
import sys
import threading
from pathlib import Path
from urllib.error import HTTPError
from urllib.parse import urlencode, urlparse
from urllib.request import HTTPRedirectHandler, Request, build_opener

MANIFEST = json.loads(Path(__file__).with_name('manifest.json').read_text())
OUTPUT_LOCK = threading.Lock()
DEVICE_LOCKS = {}
LOCK = threading.Lock()
MAX_MESSAGE = 16 * 1024 * 1024
READS = {'config': '/config', 'services': '/services', 'processes': '/processes',
         'hardware': '/plugins', 'hardware_config': '/plugin/config'}
WRITES = {**READS, 'service_control': '/services/control'}


class AdapterError(Exception):
    def __init__(self, code):
        self.code = code


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None


def request(device, endpoint, method='GET', payload=None, query=None):
    origin = device['connection']['url'].rstrip('/')
    parsed = urlparse(origin)
    if parsed.scheme not in ('http', 'https') or not parsed.hostname or parsed.username or parsed.password or parsed.path or parsed.query or parsed.fragment:
        raise AdapterError('INVALID_CONFIG')
    url = origin + '/api/v1' + endpoint
    if query:
        url += '?' + urlencode(query)
    body = None if payload is None else json.dumps(payload).encode()
    req = Request(url, data=body, method=method, headers={
        'Authorization': 'Bearer ' + device['credentials']['token'], 'Content-Type': 'application/json'})
    try:
        with build_opener(NoRedirect).open(req, timeout=5) as response:
            raw = response.read(MAX_MESSAGE + 1)
            if len(raw) > MAX_MESSAGE:
                raise AdapterError('INVALID_RESPONSE')
            value = json.loads(raw)
            if not isinstance(value, dict):
                raise AdapterError('INVALID_RESPONSE')
            return value
    except HTTPError as exc:
        raise AdapterError('AUTH_FAILED' if exc.code in (401, 403) else 'RESULT_UNKNOWN' if method == 'PUT' and exc.code >= 500 else 'DEVICE_API_ERROR') from None
    except (OSError, ValueError):
        raise AdapterError('RESULT_UNKNOWN' if method == 'PUT' else 'DEVICE_UNAVAILABLE') from None


def collect(device):
    value = request(device, '/status')
    if not isinstance(value.get('system'), dict) or not isinstance(value.get('machine'), dict):
        raise AdapterError('INVALID_RESPONSE')
    metrics = []
    units = {'cpu_percent': 'percent', 'memory_percent': 'percent', 'disk_percent': 'percent',
             'temperature_c': 'celsius', 'cpu_freq_mhz': 'megahertz',
             'memory_total_bytes': 'bytes', 'memory_used_bytes': 'bytes',
             'disk_total_bytes': 'bytes', 'disk_used_bytes': 'bytes'}
    for key, unit in units.items():
        number = value['system'].get(key)
        if isinstance(number, (int, float)) and not isinstance(number, bool):
            metrics.append({'resource_id': 'host', 'key': key, 'value': number, 'unit': unit})
    uptime = value.get('uptime_seconds')
    if isinstance(uptime, (int, float)) and not isinstance(uptime, bool):
        metrics.append({'resource_id': 'host', 'key': 'uptime_seconds', 'value': uptime, 'unit': 'seconds'})
    speed = value.get('network', {}).get('link_speed_mbps')
    if isinstance(speed, (int, float)) and not isinstance(speed, bool):
        metrics.append({'resource_id': 'host', 'key': 'network_link_speed_mbps', 'value': speed, 'unit': 'megabits_per_second'})
    for field, key in (('rx_bytes_per_second', 'network_rx_bytes_per_second'),
                       ('tx_bytes_per_second', 'network_tx_bytes_per_second')):
        rate = value.get('network', {}).get(field)
        if isinstance(rate, (int, float)) and not isinstance(rate, bool) and math.isfinite(rate) and rate >= 0:
            metrics.append({'resource_id': 'host', 'key': key, 'value': rate, 'unit': 'bytes_per_second'})
    for metric in metrics:
        source = value.get('network', {}) if metric['key'].startswith('network_') else value.get('system', {})
        at = source.get('observed_at') if metric['key'] != 'uptime_seconds' else value.get('timestamp')
        if isinstance(at, (int, float)) and not isinstance(at, bool) and math.isfinite(at) and at > 0:
            metric['observed_at'] = at
    return {'schema_version': 1, 'observed_at': value.get('timestamp'),
            'identity': {'reported_name': value['machine'].get('name'), 'category': 'host',
                         'details': value.get('device', {})},
            'collection': value.get('collection', MANIFEST['collection']),
            'capabilities': MANIFEST['capabilities'], 'metrics': metrics,
            'resources': [{'id': 'host', 'type': 'host'}],
            'alerts': value.get('display', {}).get('alerts', []),
            'alerts_complete': isinstance(value.get('display', {}).get('alerts'), list),
            'view': value}


def dispatch(method, params):
    if method in ('adapter.initialize', 'adapter.describe', 'adapter.health'):
        return MANIFEST
    device = params['device']
    with LOCK:
        device_lock = DEVICE_LOCKS.setdefault(device['id'], threading.Lock())
    with device_lock:
        if method == 'device.validate':
            return request(device, '/health')
        if method == 'device.collect':
            return collect(device)
        if method == 'device.logs.query':
            return request(device, '/logs', query=params.get('query'))
        if method == 'device.read':
            return request(device, READS[params['resource']], query=params.get('query'))
        if method == 'device.write':
            return request(device, WRITES[params['resource']], 'PUT', params.get('payload'))
        raise AdapterError('UNSUPPORTED')


def handle(message):
    try:
        result = dispatch(message['method'], message.get('params', {}))
        reply = {'jsonrpc': '2.0', 'id': message['id'], 'result': result}
    except AdapterError as exc:
        reply = {'jsonrpc': '2.0', 'id': message['id'], 'error': {'code': -32000, 'data': exc.code}}
    except Exception:
        # Do not copy vendor exception messages: they may contain credentials or URLs.
        reply = {'jsonrpc': '2.0', 'id': message['id'], 'error': {'code': -32000, 'data': 'DEVICE_UNAVAILABLE'}}
    raw = json.dumps(reply, allow_nan=False, separators=(',', ':'))
    with OUTPUT_LOCK:
        print(raw, flush=True)


def main():
    # Bound both running and queued calls even if the parent violates its limits.
    slots = threading.BoundedSemaphore(8)
    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
        while True:
            line = sys.stdin.buffer.readline(MAX_MESSAGE + 1)
            if not line:
                break
            if len(line) > MAX_MESSAGE or not line.endswith(b'\n'):
                break
            message = json.loads(line)
            slots.acquire()
            task = pool.submit(handle, message)
            task.add_done_callback(lambda _: slots.release())


if __name__ == '__main__':
    main()