path: root/pyagl/services/base.py

# Copyright (C) 2020 Konsulko Group
# Author: Edi Feschiyan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from subprocess import check_output
from json import JSONDecodeError
from parse import Result, parse
from websockets import connect
from random import randint
from enum import IntEnum
from typing import Union
import asyncssh
from asyncssh import SSHClientConnection
import argparse
import asyncio
import binascii
import logging
import json
import sys
import os
import re
# logging.getLogger('AGLBaseService')
# logging.basicConfig(level=logging.DEBUG)


# AFB message type
class AFBT(IntEnum):
    REQUEST = 2,
    RESPONSE = 3,
    ERROR = 4,
    EVENT = 5


msgq = {}
# AppFrameworkBinder responses/events have 3 items in the list - [messagetype, sessionid, payload]
AFBLEN = 3


def newrand():
    while True:
        bs = os.urandom(5)
        result = bs[0] * bs[1] * bs[2] * bs[3] + bs[4]
        yield result


def addrequest(msgid, msg):
    msgq[msgid] = {'request': msg, 'response': None}


def addresponse(msgid, msg):
    if msgid in msgq.keys():
        msgq[msgid]['response'] = msg


class AFBResponse:
    type: AFBT
    msgid: int
    data: dict
    api: str
    status: str
    info: str

    def __init__(self, data: list):
        if not isinstance(data, list):
            raise TypeError('data parameter is not list type')
        if type(data[0]) is not int:
            logging.debug(f'Received a response with non-integer message type {binascii.hexlify(data[0])}')
            raise ValueError('Received a response with non-integer message type')
        if data[0] not in AFBT._value2member_map_:
            raise ValueError(f'Received a response with invalid message type {data[0]}')
        self.type = AFBT(data[0])

        if self.type == AFBT.RESPONSE:
            if 'request' not in data[2]:
                logging.error(f'Received malformed or invalid response without "request" dict - {data}')
            if not str.isnumeric(data[1]):
                raise ValueError(f'Received a response with non-numeric message id {data[1]}')
            else:
                self.msgid = int(data[1])
            self.status = data[2]['request'].get('status')
            self.info = data[2]['request'].get('info')
            self.data = data[2].get('response')

        elif self.type == AFBT.EVENT:
            self.api = data[1]
            self.data = data[2].get('data')

        elif self.type == AFBT.ERROR:
            logging.debug(f'AFB returned erroneous response {data}')
            self.msgid = int(data[1])
            self.status = data[2]['request'].get('status')
            self.info = data[2]['request'].get('info')
            self.data = data[2]['request'].get('data')

    def __str__(self):  # for debugging purposes
        if self.type == AFBT.EVENT:
            return f'[{self.type.name}][{self.api}][Data: {self.data}]'
        else:
            return f'[{self.type.name}][Status: {self.status}][{self.msgid}][Info: {self.info}][Data: {self.data}]'


class AGLBaseService:
    api: str
    url: str
    ip : str
    port = None
    token: str
    uuid: str
    service = None
    logger = None

    @staticmethod
    def getparser():
        parser = argparse.ArgumentParser(description='Utility to interact with agl-service-* via it\'s websocket')
        parser.add_argument('-l', '--loglevel', help='Level of logging verbosity', default='INFO',
                            choices=list(logging._nameToLevel.keys()))
        parser.add_argument('ipaddr', default=os.environ.get('AGL_TGT_IP', 'localhost'), help='AGL host address')
        parser.add_argument('--port', default=os.environ.get('AGL_TGT_PORT', None), help=f'AGL service websocket port')
        parser.add_argument('--listener', default=False, help='Register a listener for incoming events', action='store_true')
        parser.add_argument('--subscribe', type=str, help='Subscribe to event type', action='append', metavar='event')
        parser.add_argument('--unsubscribe', type=str, help='Unsubscribe from event type', action='append', metavar='event')
        # parser.add_argument('--json', type=str, help='Send your own json string')
        # parser.add_argument('--verb', type=str, help='Send the json above to specific verb')
        # parser.add_argument('--api', type=str, help='Send the above two to a specific api')
        return parser

    def __init__(self, api: str, ip: str, port: str = None, url: str = None,
                 token: str = 'HELLO', uuid: str = 'magic', service: str = None,
                 runservice: bool = False, timeout: float = 5.0):
        self.api = api
        self.url = url
        self.ip = ip
        self.port = port
        self.token = token
        self.uuid = uuid
        self.service = service
        self.runsvc = runservice
        self.logger = logging.getLogger(service)
        try:
            self.timeout = float(os.environ.get('AGL_TEST_TIMEOUT', timeout))
        except ValueError:
            self.timeout = 5.0

    def __await__(self):
        return self._async_init().__await__()

    async def __aenter__(self):
        return self._async_init()

    async def _async_init(self):
        # setting ping_timeout to None because AFB does not support websocket ping
        # if set to !None, the library will close the socket after the default timeout
        if self.port is None:
            serviceport = await self.portfinder(runservice=self.runsvc)
            if serviceport is not None:
                self.port = serviceport
            else:
                print("Service port: ", serviceport)
                self.logger.error('Unable to find port')
                exit(1)

        URL = f'ws://{self.ip}:{self.port}/api?token={self.token}&uuid={self.uuid}'
        self._conn = connect(close_timeout=0, uri=URL, subprotocols=['x-afb-ws-json1'], ping_timeout=None, compression=None)
        self.websocket = await self._conn.__aenter__()
        return self

    async def __aexit__(self, *args, **kwargs):
        await self._conn.__aexit__(*args, **kwargs)

    async def close(self):
        await self._conn.__aexit__(*sys.exc_info())

    async def send(self, message):
        await self.websocket.send(message)

    async def receive(self):
        return await self.websocket.recv()

    async def portfinder(self, runservice=False):
        fieldsstr = '{sl}: {local_address} {rem_address} {st} {tx_queue}:{rx_queue} {tr}:{tmwhen} {retrnsmt} {uid}' \
                    ' {timeout} {inode} {sref_cnt} {memloc} {rto} {pred_sclk} {ackquick} {congest} {slowstart}'
        # TODO: handle ssh timeouts, asyncssh does not support it apparently,
        #       and connect returns context_manager which cannot be used with
        #       asyncio.wait_for
        if self.ip == 'localhost' or self.ip == '127.0.0.1':
            servicename = check_output(f"systemctl --all --no-legend | grep {self.service}-- | awk '{{print $1}}'", shell=True)
            servicename = servicename.decode().strip().splitlines()
            # it is possible that two matching services are returned
            # - one for uid 0 as root and 1001 as agl-driver, we want the latter
            servicename = servicename[-1] if len(servicename) else ''
            if self.service not in servicename:
                if runservice:
                    await self.runafmservice()
                    servicename = check_output(f"systemctl --all --no-legend | grep {self.service}-- | awk '{{print $1}}'",
                                              shell=True)  # retry getting service name after starting
                    servicename = servicename.decode().strip()
                else:
                    self.logger.error(f"Service matching pattern - '{self.service}' - NOT FOUND")
                    exit(1)
            pid = check_output(f'systemctl show --property MainPID --value {servicename}'.encode(), shell=True)
            pid = int(pid.decode().strip(), 10)
            if pid == 0 and servicename != '':
                self.logger.warning(f'Service {servicename} is stopped')
                if runservice:
                    self.logger.warning(f' Trying to start service {servicename}')
                    await self.startsystemdservice(servicename=servicename)
                    pid = check_output(f'systemctl show --property MainPID --value {servicename}'.encode(), shell=True)
                    pid = int(pid.decode().strip(), 10)

                else:
                    return None
            else:
                self.logger.debug(f'Service PID: {str(pid)}')

            sockets = check_output(f'find /proc/{pid}/fd/ | xargs readlink | grep socket'.encode(), shell=True)
            sockets = sockets.decode().strip()
            inodes = frozenset(re.findall("socket:\\[(.*)\\]", sockets))
            self.logger.debug(f"Socket inodes: {inodes}")

            procnettcp = check_output('cat /proc/net/tcp'.encode(), shell=True)
            procnettcp = procnettcp.decode().splitlines()[1:]

        else:
            async with asyncssh.connect(self.ip, username='root', known_hosts=None) as c:
                servicename = await c.run(f"systemctl --all --no-legend | grep {self.service}-- | awk '{{print $1}}'", check=False)
                servicename = servicename.stdout.strip().splitlines()
                print(f"matching services: {servicename}")
                servicename = servicename[-1] if len(servicename) else ''
                # it is possible that two matching services
                # are returned - one for uid 0 as root and 1001 as agl-driver, we want the latter
                if self.service not in servicename:
                    if runservice:
                        await self.runafmservice(connection=c)
                        servicename = await c.run(f"systemctl --all --no-legend | grep {self.service}-- | awk '{{print $1}}'",
                                                  check=False)  # retry getting service name after starting
                        servicename = servicename.stdout.strip()
                    else:
                        self.logger.error(f"Service matching pattern - '{self.service}' - NOT FOUND")
                        exit(1)
                pidres = await c.run(f'systemctl show --property MainPID --value {servicename}')
                pid = int(pidres.stdout.strip(), 10)
                if pid == 0 and servicename != '':
                    self.logger.warning(f'Service {servicename} is stopped')
                    if runservice:
                        self.logger.warning(f' Trying to start service {servicename}')
                        await self.startsystemdservice(servicename=servicename, connection=c)
                        pidres = await c.run(f'systemctl show --property MainPID --value {servicename}')
                        pid = int(pidres.stdout.strip(), 10)
                    else:
                        return None
                else:
                    self.logger.debug(f'Service PID: {str(pid)}')

                sockets = await c.run(f'find /proc/{pid}/fd/ | xargs readlink | grep socket')
                inodes = frozenset(re.findall("socket:\\[(.*)\\]", sockets.stdout))
                self.logger.debug(f"Socket inodes: {inodes}")

                procnettcp = await c.run('cat /proc/net/tcp')
                procnettcp = procnettcp.stdout.splitlines()[1:]

        self.logger.debug(procnettcp)
        tcpsockets = [' '.join(l.split()) for l in procnettcp]
        # different lines with less stats appear sometimes, parse will return None, so ignore 'None' lines
        parsedtcpsockets = []
        for l in tcpsockets:
            res = parse(fieldsstr, l)
            if isinstance(res, Result):
                parsedtcpsockets.append(res)

        self.logger.debug(parsedtcpsockets)
        socketinodesbythisprocess = [l for l in parsedtcpsockets if
                                     isinstance(l, Result) and
                                     l.named['inode'] in inodes and
                                     # 0A is listening state for the socket
                                     l.named['st'] == '0A']

        self.logger.debug(socketinodesbythisprocess)
        for s in socketinodesbythisprocess:
            _, port = tuple(parse('{}:{}', s['local_address']))
            port = int(port, 16)
            if port >= 30000:  # the port is above 30000 range, 8080 is some kind of proxy
                self.logger.debug(f'Service running at port {port}')
                return port

    async def runafmservice(self, servicename: str = None, connection: SSHClientConnection = None):
        name = servicename if servicename is not None else self.service
        pid = None
        result = None
        if self.ip == '127.0.0.1' or self.ip == 'localhost':  # running locally, ssh not needed
            result = check_output(f'afm-util start {name}', shell=True).decode().strip()
        else:  # running remotely
            if connection is not None:
                result = await connection.run(f'afm-util start {self.service}')
                result = result.stdout.strip()
            else:
                self.logger.error('Trying to start service remotely but no SSHClientConnection given')

        if result.isnumeric:
            self.logger.debug(f'Started service via afm-util, PID: {result}')
            pid = int(result)
        else:
            self.logger.error(f'Unable to start service via afm-util: {result}')

        return pid

    @staticmethod
    async def startsystemdservice(servicename: str, connection: SSHClientConnection = None):
        result = None
        if connection is not None:
            result = await connection.run(f'systemctl start {servicename}')
            result = result.stdout.strip()
        else:
            result = check_output(f'systemctl start {servicename}', shell=True)
            print(result)

    async def listener(self, stdout: bool = False, timeout: float = None):
        while True:
            raw = None
            try:
                if timeout is None:
                    timeout = self.timeout
                raw = await asyncio.wait_for(self.response(), timeout)
                data = AFBResponse(raw)
                if stdout:
                    print(data)
                yield data

            except asyncio.TimeoutError:
                self.logger.warning("Response wait timed out")
                yield None

    async def response(self):
        try:
            msg = await asyncio.wait_for(self.websocket.recv(), self.timeout)
            try:
                data = json.loads(msg)
                self.logger.debug('[AGL] -> ' + msg)
                if isinstance(data, list):
                    # check whether the received response is an answer to previous query and queue it for debugging
                    if len(data) == AFBLEN and data[0] == AFBT.RESPONSE and str.isnumeric(data[1]):
                        msgid = int(data[1])
                        if msgid in msgq:
                            addresponse(msgid, data)
                    return data
            except JSONDecodeError:
                self.logger.warning("Not decoding a non-json message")

        except KeyboardInterrupt:
            self.logger.debug("Received keyboard interrupt, exiting")
        except asyncio.CancelledError:
            self.logger.warning("Websocket listener coroutine stopped")
        except asyncio.TimeoutError:
            self.logger.warning("Response wait timed out")
        except Exception as e:
            self.logger.error("Unhandled seal: " + str(e))

    async def afbresponse(self):
        resp = None
        try:
            resp = await self.response()
            resp = AFBResponse(resp)
        except TypeError as v:  # on asyncio.timeout()-ed response will return None and AFBResponse() will
            # raise an error, catch it. probably more appropriate to handle it in AFBResponse()
            self.logger.warning("Something probably timedout on response(), just tried to parse None to AFBResponse()")

        return resp

    async def request(self, verb: str, values: Union[str, dict] = "", msgid: int = None, api=None):
        msgid = next(newrand()) if msgid is None else msgid
        if api is not None:
            l = json.dumps([AFBT.REQUEST, str(msgid), f'{api}/{verb}', values])
        else:
            l = json.dumps([AFBT.REQUEST, str(msgid), f'{self.api}/{verb}', values])
        self.logger.debug(f'[AGL] <- {l}')
        await self.send(l)
        return msgid

    async def subscribe(self, event):
        return await self.request('subscribe', {'value': f'{event}'})  # some services may use 'event' instead 'value'

    async def unsubscribe(self, event):
        return await self.request('unsubscribe', {'value': f'{event}'})