esp-idf/examples/protocols/websocket/example_test.py

from __future__ import print_function
from __future__ import unicode_literals
import re
import os
import socket
import hashlib
import base64
from threading import Thread

import ttfw_idf


def get_my_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        # doesn't even have to be reachable
        s.connect(('10.255.255.255', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP


# Simple Websocket server for testing purposes
class Websocket:
    HEADER_LEN = 6

    def __init__(self, port):
        self.port = port
        self.socket = socket.socket()
        self.socket.settimeout(10.0)

    def __enter__(self):
        try:
            self.socket.bind(('', self.port))
        except socket.error as e:
            print("Bind failed:{}".format(e))
            raise

        self.socket.listen(1)
        self.server_thread = Thread(target=self.run_server)
        self.server_thread.start()

    def __exit__(self, exc_type, exc_value, traceback):
        self.server_thread.join()
        self.socket.close()
        self.conn.close()

    def run_server(self):
        self.conn, address = self.socket.accept()  # accept new connection
        self.conn.settimeout(10.0)
        print("Connection from: {}".format(address))

        self.establish_connection()

        # Echo data until client closes connection
        self.echo_data()

    def establish_connection(self):
        while True:
            try:
                # receive data stream. it won't accept data packet greater than 1024 bytes
                data = self.conn.recv(1024).decode()
                if not data:
                    # exit if data is not received
                    raise

                if "Upgrade: websocket" in data and "Connection: Upgrade" in data:
                    self.handshake(data)
                    return
            except socket.error as err:
                print("Unable to establish a websocket connection: {}, {}".format(err))
                raise

    def handshake(self, data):
        # Magic string from RFC
        MAGIC_STRING = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
        headers = data.split("\r\n")

        for header in headers:
            if "Sec-WebSocket-Key" in header:
                client_key = header.split()[1]

        if client_key:
            resp_key = client_key + MAGIC_STRING
            resp_key = base64.standard_b64encode(hashlib.sha1(resp_key.encode()).digest())

            resp = "HTTP/1.1 101 Switching Protocols\r\n" + \
                "Upgrade: websocket\r\n" + \
                "Connection: Upgrade\r\n" + \
                "Sec-WebSocket-Accept: {}\r\n\r\n".format(resp_key.decode())

            self.conn.send(resp.encode())

    def echo_data(self):
        while(True):
            try:
                header = bytearray(self.conn.recv(self.HEADER_LEN, socket.MSG_WAITALL))
                if not header:
                    # exit if data is not received
                    return

                # Remove mask bit
                payload_len = ~(1 << 7) & header[1]

                payload = bytearray(self.conn.recv(payload_len, socket.MSG_WAITALL))
                frame = header + payload

                decoded_payload = self.decode_frame(frame)

                echo_frame = self.encode_frame(decoded_payload)
                self.conn.send(echo_frame)
            except socket.error as err:
                print("Stopped echoing data: {}".format(err))

    def decode_frame(self, frame):
        # Mask out MASK bit from payload length, this len is only valid for short messages (<126)
        payload_len = ~(1 << 7) & frame[1]

        mask = frame[2:self.HEADER_LEN]

        encrypted_payload = frame[self.HEADER_LEN:self.HEADER_LEN + payload_len]
        payload = bytearray()

        for i in range(payload_len):
            payload.append(encrypted_payload[i] ^ mask[i % 4])

        return payload

    def encode_frame(self, payload):
        # Set FIN = 1 and OP_CODE = 1 (text)
        header = (1 << 7) | (1 << 0)

        frame = bytearray(header)
        frame.append(len(payload))
        frame += payload

        return frame


def test_echo(dut):
    dut.expect("WEBSOCKET_EVENT_CONNECTED")
    for i in range(0, 10):
        dut.expect(re.compile(r"Received=hello (\d)"))
    dut.expect("Websocket Stopped")


@ttfw_idf.idf_example_test(env_tag="Example_WIFI")
def test_examples_protocol_websocket(env, extra_data):
    """
    steps:
      1. join AP
      2. connect to uri specified in the config
      3. send and receive data
    """
    dut1 = env.get_dut("websocket", "examples/protocols/websocket")
    # check and log bin size
    binary_file = os.path.join(dut1.app.binary_path, "websocket-example.bin")
    bin_size = os.path.getsize(binary_file)
    ttfw_idf.log_performance("websocket_bin_size", "{}KB".format(bin_size // 1024))
    ttfw_idf.check_performance("websocket_bin_size", bin_size // 1024)

    try:
        if "CONFIG_WEBSOCKET_URI_FROM_STDIN" in dut1.app.get_sdkconfig():
            uri_from_stdin = True
        else:
            uri = dut1.app.get_sdkconfig()["CONFIG_WEBSOCKET_URI"].strip('"')
            uri_from_stdin = False

    except Exception:
        print('ENV_TEST_FAILURE: Cannot find uri settings in sdkconfig')
        raise

    # start test
    dut1.start_app()

    if uri_from_stdin:
        server_port = 4455
        with Websocket(server_port):
            uri = "ws://{}:{}".format(get_my_ip(), server_port)
            print("DUT connecting to {}".format(uri))
            dut1.expect("Please enter uri of websocket endpoint", timeout=30)
            dut1.write(uri)
            test_echo(dut1)

    else:
        print("DUT connecting to {}".format(uri))
        test_echo(dut1)


if __name__ == '__main__':
    test_examples_protocol_websocket()