1# Copyright 2015-2021 Espressif Systems (Shanghai) CO LTD
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import json
16import time
17
18from .output_helpers import red_print, yellow_print
19
20try:
21    import websocket
22except ImportError:
23    # This is needed for IDE integration only.
24    pass
25
26
27class WebSocketClient(object):
28    """
29    WebSocket client used to advertise debug events to WebSocket server by sending and receiving JSON-serialized
30    dictionaries.
31
32    Advertisement of debug event:
33    {'event': 'gdb_stub', 'port': '/dev/ttyUSB1', 'prog': 'build/elf_file'} for GDB Stub, or
34    {'event': 'coredump', 'file': '/tmp/xy', 'prog': 'build/elf_file'} for coredump,
35    where 'port' is the port for the connected device, 'prog' is the full path to the ELF file and 'file' is the
36    generated coredump file.
37
38    Expected end of external debugging:
39    {'event': 'debug_finished'}
40    """
41
42    RETRIES = 3
43    CONNECTION_RETRY_DELAY = 1
44
45    def __init__(self, url):  # type: (str) -> None
46        self.url = url
47        self._connect()
48
49    def _connect(self):  # type: () -> None
50        """
51        Connect to WebSocket server at url
52        """
53
54        self.close()
55        for _ in range(self.RETRIES):
56            try:
57                self.ws = websocket.create_connection(self.url)
58                break  # success
59            except NameError:
60                raise RuntimeError('Please install the websocket_client package for IDE integration!')
61            except Exception as e:  # noqa
62                red_print('WebSocket connection error: {}'.format(e))
63            time.sleep(self.CONNECTION_RETRY_DELAY)
64        else:
65            raise RuntimeError('Cannot connect to WebSocket server')
66
67    def close(self):  # type: () -> None
68        try:
69            self.ws.close()
70        except AttributeError:
71            # Not yet connected
72            pass
73        except Exception as e:  # noqa
74            red_print('WebSocket close error: {}'.format(e))
75
76    def send(self, payload_dict):  # type: (dict) -> None
77        """
78        Serialize payload_dict in JSON format and send it to the server
79        """
80        for _ in range(self.RETRIES):
81            try:
82                self.ws.send(json.dumps(payload_dict))
83                yellow_print('WebSocket sent: {}'.format(payload_dict))
84                break
85            except Exception as e:  # noqa
86                red_print('WebSocket send error: {}'.format(e))
87                self._connect()
88        else:
89            raise RuntimeError('Cannot send to WebSocket server')
90
91    def wait(self, expect_iterable):  # type: (list) -> None
92        """
93        Wait until a dictionary in JSON format is received from the server with all (key, value) tuples from
94        expect_iterable.
95        """
96        for _ in range(self.RETRIES):
97            try:
98                r = self.ws.recv()
99            except Exception as e:
100                red_print('WebSocket receive error: {}'.format(e))
101                self._connect()
102                continue
103            obj = json.loads(r)
104            if all([k in obj and obj[k] == v for k, v in expect_iterable]):
105                yellow_print('WebSocket received: {}'.format(obj))
106                break
107            red_print('WebSocket expected: {}, received: {}'.format(dict(expect_iterable), obj))
108        else:
109            raise RuntimeError('Cannot receive from WebSocket server')
110