openrvdas.logger.readers.sealog_reader

  1#!/usr/bin/env python3
  2
  3import asyncio
  4import logging
  5import queue
  6import ssl
  7import json
  8import threading
  9import time
 10import websockets
 11
 12from logger.utils.das_record import DASRecord  # noqa: E402
 13from logger.utils.timestamp import timestamp  # noqa: E402
 14
 15
 16################################################################################
 17class SealogReader():
 18    """Connect to a websocket served by a WebsocketWriter, and service read()
 19    requests from it.
 20    """
 21
 22    def __init__(self, uri, client_wsid=None, subs=None, check_cert=False, **kwargs):
 23        """
 24        ```
 25        uri - Hostname, port and protocol, (e.g. wss://localhost:8080) of the
 26              target Sealog Server
 27
 28        client_wsid - Unique identifier, useful when running multiple
 29                      SealogReaders
 30
 31        subs - Websocket Pub/Subs to subscribe to when listening to the Sealog
 32               Server
 33
 34        check_cert - If True, and uri protocol is 'wss', check the server's TLS
 35                     certificate for validity; if a str, use as local filepath
 36                     location of .pem file to check against.
 37        ```
 38        """
 39        super().__init__(**kwargs)
 40
 41        self.uri = uri
 42        self.client_wsid = client_wsid or 'OpenRVDAS_Sealog_Reader'
 43        self.subs = subs or ['/ws/status/newEvents', '/ws/status/updateEvents']
 44        self.check_cert = check_cert
 45
 46        self.ping = {
 47            'type': 'ping',
 48            'id': self.client_wsid
 49        }
 50
 51        self.hello = {
 52            'type': 'hello',
 53            'id': self.client_wsid,
 54            'version': '2',
 55            'subs': self.subs
 56        }
 57
 58        # We won't initialize our websocket connection until the first read()
 59        # call. At that point we'll launch an async process in a separate
 60        # thread that will wait for data from the websocket and put it in a
 61        # queue that read() will pop from.
 62        self.websocket_thread = None
 63        self.queue = queue.Queue()
 64        self.quit_flag = False
 65
 66    ############################
 67    def _start_websocket(self):
 68        """We'll run this in a separate thread as soon as we get our first
 69        call to read()."""
 70
 71        ############################
 72        def _event_to_das_record(event) -> DASRecord:
 73
 74            fields = {
 75                'event_id': event.get('id'),
 76                'event_value': event.get('event_value'),
 77                'event_author': event.get('event_author'),
 78                'event_free_text': event.get('event_free_text'),
 79                **{
 80                    f"event_option_{option['event_option_name']}": option["event_option_value"]
 81                    for option in event.get("event_options", [])
 82                }
 83            }
 84
 85            return DASRecord(
 86                data_id=self.client_wsid,
 87                message_type='sealog_event',
 88                timestamp=timestamp(event.get('ts')),
 89                fields=fields
 90            )
 91
 92        ############################
 93        async def _websocket_loop(self):
 94            """Asynchronous inner function that will connect to websocket,
 95            iteratively try to read from it and put the result in our queue.
 96            """
 97            # Iterate if we lose the websocket for some reason other than a 'quit'
 98            while not self.quit_flag:
 99                try:
100                    if self.uri.find('wss://') == 0:  # using wss
101                        # If check_cert is a str, take it as the location of the
102                        # .pem file we'll check for validity. Otherwise, if not
103                        # False, take as a bool to verify by own means.
104                        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
105                        if self.check_cert:
106                            if isinstance(self.check_cert, str):
107                                ssl_context.load_verify_locations(self.check_cert)
108                            else:
109                                ssl_context.verify_mode = ssl.CERT_REQUIRED
110                        else:
111                            ssl_context.check_hostname = False
112                            ssl_context.verify_mode = ssl.CERT_NONE
113
114                    else:  # not using wss
115                        ssl_context = None
116
117                    logging.debug(f'SealogReader connecting to {self.uri}')
118                    async with websockets.connect(self.uri, ssl=ssl_context) as ws:  # type: ignore
119
120                        await ws.send(json.dumps(self.hello))
121                        logging.info(f'Connected to SealogReader at {self.uri}')
122
123                        while not self.quit_flag:
124                            msg = await ws.recv()
125                            msg_obj = json.loads(msg)
126
127                            if msg_obj['type'] and msg_obj['type'] == 'ping':
128                                await ws.send(json.dumps(self.ping))
129                                continue
130
131                            if msg_obj['type'] and msg_obj['type'] == 'pub':
132
133                                event = msg_obj['message']
134                                logging.debug(
135                                    f'SealogReader got record {json.dumps(event, indent=2)}'
136                                )
137                                self.queue.put(_event_to_das_record(event))
138
139                except BrokenPipeError:
140                    logging.info(f'SealogReader BrokenPipeError connecting to {self.uri}')
141                except AttributeError as e:
142                    logging.warning(f'SealogReader websocket loop error: {e}')
143                except websockets.exceptions.ConnectionClosed:  # type: ignore
144                    logging.info('SealogReader lost websocket connection to '
145                                 'data server; trying to reconnect.')
146                    await asyncio.sleep(0.2)
147                except OSError as e:
148                    logging.info('Unable to connect to websocket. Sleeping to try again...')
149                    logging.info(f'Connection error: {e}')
150                    await asyncio.sleep(2)
151
152        # In the outer function, get a new event loop and fire up the
153        # inner, async routine.
154        # Could we also use asyncio.ensure_future(_websocket_loop(self)) ?
155        websocket_event_loop = asyncio.new_event_loop()
156        websocket_event_loop.run_until_complete(_websocket_loop(self))
157        websocket_event_loop.close()
158
159    ############################
160    def quit(self, seconds=0):
161        """Sleep N seconds, then signal quit."""
162        time.sleep(seconds)
163        self.quit_flag = True
164
165    ############################
166    def read(self):
167        """Read/wait for data from the websocket."""
168
169        # If we've not yet fired up the websocket thread, do that now.
170        if not self.websocket_thread:
171            self.websocket_thread = threading.Thread(
172                name='websocket_thread',
173                target=self._start_websocket,
174                daemon=True)
175            self.websocket_thread.start()
176
177        # Use a timeout in our queue get() so we can periodically check if
178        # we've gotten a 'quit'
179        while not self.quit_flag:
180            try:
181                result = self.queue.get(block=True, timeout=2)
182                logging.debug('Got result from queue: %s', result)
183                return result
184            except queue.Empty:
185                logging.debug('get() timed out - trying again')
186
187        # If we've fallen out because of a quit...
188        return None