Skip to content

Client

SusClient

This class is responsible for managing the client.

Source code in sus/client/client.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class SusClient:
    """
    This class is responsible for managing the client.
    """
    protocol: SusClientProtocol

    def __init__(self, addr: tuple[str, int], ppks: str, protocol_id: bytes):
        """
        Initializes the client.
        :param addr: Server address
        :param ppks: Server public key
        :param protocol_id:  Protocol ID (any bytestring)
        """
        self.server_addr = addr
        self.ppks = X25519PublicKey.from_public_bytes(bytes.fromhex(ppks))
        self.protocol_id = protocol_id

        self.logger = logging.getLogger(f"sussus")

    def __del__(self):
        self.disconnect()

    @property
    def connected(self):
        """
        True if the client is connected to the server.
        """
        return hasattr(self, "protocol") and self.protocol.state == ConnectionState.CONNECTED

    async def start(self, handlers: Iterable[MessageHandler] = None):
        """
        This coroutine is responsible for starting the client. Blocks until the client is connected.
        It also registers message handlers, called when a message is received.
        :param handlers:
        :return:
        """
        await self.connect()
        for handler in handlers or []:
            self.protocol.add_message_handler(handler)

    def __key_exchange(self, epks_ns_port: bytes, wallet: Wallet) -> Wallet:
        """
        This function is responsible for performing the key exchange.
        :param epks_ns_port: received (epks, ns, port) from server
        :param wallet: wallet containing the client's keys
        :return: wallet containing the shared secret
        """

        if len(epks_ns_port) != 40:
            raise MalformedPacket("Invalid key response length")
        # 4. receive (epks, ns, port) from server
        wallet.epks = X25519PublicKey.from_public_bytes(epks_ns_port[:32])
        wallet.ns = epks_ns_port[32:40]
        self.logger.info("received keys, starting handshake")
        # 5. compute ecps = X25519(eskc, ppks)
        ecps = wallet.eskc.exchange(wallet.ppks)
        eces = wallet.eskc.exchange(wallet.epks)
        # 6. compute key = H(eces, ecps, nc, ns, ppks, epks, epkc)
        wallet.shared_secret = blake3(
            eces + ecps + wallet.nc + wallet.ns +
            wallet.ppks.public_bytes(Encoding.Raw, PublicFormat.Raw) +
            wallet.epks.public_bytes(Encoding.Raw, PublicFormat.Raw) +
            wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw)).digest()
        self.logger.info("shared secret: %s", wallet.shared_secret.hex())

        # 7. compute token = H(epkc, epks, nc, ns)
        self.logger.info("\n".join([
            f"epkc: {wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw).hex()}",
            f"epks: {wallet.epks.public_bytes(Encoding.Raw, PublicFormat.Raw).hex()}",
            f"nc: {wallet.nc.hex()}",
            f"ns: {wallet.ns.hex()}"
        ]))
        wallet.token = blake3(wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw) +
                              wallet.epks.public_bytes(Encoding.Raw, PublicFormat.Raw) +
                              wallet.nc + wallet.ns).digest()
        return wallet

    async def connect(self):
        """
        This coroutine is responsible for connecting to the server.
        Performs the key exchange and starts the handshake.
        """
        self.logger.info(f"connecting to server ({self.server_addr[0]}:{self.server_addr[1]})")

        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.connect(self.server_addr)
        sock.setblocking(False)
        sock.settimeout(5)

        eskc = X25519PrivateKey.generate()
        epkc = eskc.public_key()
        nc = urandom(8)
        wallet = Wallet(ppks=self.ppks, eskc=eskc, epkc=epkc, nc=nc)

        # try:
        sock.send(wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw) + wallet.nc)
        data = sock.recv(40)
        # except (ConnectionError, TimeoutError):
        #     self.logger.error("failed to connect to server")
        # return False

        wallet = self.__key_exchange(data, wallet)

        self.logger.info("received keys, starting handshake")

        _, self.protocol = await asyncio.get_event_loop().create_datagram_endpoint(
            lambda: SusClientProtocol(wallet, self.protocol_id),
            sock=sock
        )
        await self.protocol.handshake_event.wait()
        # return True

    def send(self, data: bytes):
        """
        Sends a message to the server.
        :param data: message to send as bytes
        """
        if not self.protocol:
            self.logger.warning("not connected to server")
            return
        self.protocol.send(data)

    def disconnect(self):
        """
        Disconnects from the server.
        """
        if not hasattr(self, "protocol"):
            self.logger.warning("not connected to server")
            return
        try:
            asyncio.get_running_loop()
            self.protocol.disconnect()
        except RuntimeError:  # not running in event loop
            pass
        self.logger.info(f"disconnected from server ({self.server_addr[0]}:{self.server_addr[1]})")

    async def keep_alive(self):
        """
        Convenience coroutine that waits until the client is disconnected.
        """
        if not hasattr(self, "protocol"):
            self.logger.warning("not connected to server")
            return
        try:
            await self.protocol.diconnection_event.wait()
        except asyncio.CancelledError:
            self.logger.info("exiting...")

connected property

True if the client is connected to the server.

__init__(addr, ppks, protocol_id)

Initializes the client.

Parameters:

Name Type Description Default
addr tuple[str, int]

Server address

 required
ppks str

Server public key

 required
protocol_id bytes

Protocol ID (any bytestring)

 required
Source code in sus/client/client.py 
26
27
28
29
30
31
32
33
34
35
36
37
def __init__(self, addr: tuple[str, int], ppks: str, protocol_id: bytes):
    """
    Initializes the client.
    :param addr: Server address
    :param ppks: Server public key
    :param protocol_id:  Protocol ID (any bytestring)
    """
    self.server_addr = addr
    self.ppks = X25519PublicKey.from_public_bytes(bytes.fromhex(ppks))
    self.protocol_id = protocol_id

    self.logger = logging.getLogger(f"sussus")

__key_exchange(epks_ns_port, wallet)

This function is responsible for performing the key exchange.

Parameters:

Name Type Description Default
epks_ns_port bytes

received (epks, ns, port) from server

 required
wallet Wallet

wallet containing the client's keys

 required

Returns:

Type Description
Wallet

wallet containing the shared secret
Source code in sus/client/client.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def __key_exchange(self, epks_ns_port: bytes, wallet: Wallet) -> Wallet:
    """
    This function is responsible for performing the key exchange.
    :param epks_ns_port: received (epks, ns, port) from server
    :param wallet: wallet containing the client's keys
    :return: wallet containing the shared secret
    """

    if len(epks_ns_port) != 40:
        raise MalformedPacket("Invalid key response length")
    # 4. receive (epks, ns, port) from server
    wallet.epks = X25519PublicKey.from_public_bytes(epks_ns_port[:32])
    wallet.ns = epks_ns_port[32:40]
    self.logger.info("received keys, starting handshake")
    # 5. compute ecps = X25519(eskc, ppks)
    ecps = wallet.eskc.exchange(wallet.ppks)
    eces = wallet.eskc.exchange(wallet.epks)
    # 6. compute key = H(eces, ecps, nc, ns, ppks, epks, epkc)
    wallet.shared_secret = blake3(
        eces + ecps + wallet.nc + wallet.ns +
        wallet.ppks.public_bytes(Encoding.Raw, PublicFormat.Raw) +
        wallet.epks.public_bytes(Encoding.Raw, PublicFormat.Raw) +
        wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw)).digest()
    self.logger.info("shared secret: %s", wallet.shared_secret.hex())

    # 7. compute token = H(epkc, epks, nc, ns)
    self.logger.info("\n".join([
        f"epkc: {wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw).hex()}",
        f"epks: {wallet.epks.public_bytes(Encoding.Raw, PublicFormat.Raw).hex()}",
        f"nc: {wallet.nc.hex()}",
        f"ns: {wallet.ns.hex()}"
    ]))
    wallet.token = blake3(wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw) +
                          wallet.epks.public_bytes(Encoding.Raw, PublicFormat.Raw) +
                          wallet.nc + wallet.ns).digest()
    return wallet

connect() async

This coroutine is responsible for connecting to the server. Performs the key exchange and starts the handshake.

Source code in sus/client/client.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
async def connect(self):
    """
    This coroutine is responsible for connecting to the server.
    Performs the key exchange and starts the handshake.
    """
    self.logger.info(f"connecting to server ({self.server_addr[0]}:{self.server_addr[1]})")

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.connect(self.server_addr)
    sock.setblocking(False)
    sock.settimeout(5)

    eskc = X25519PrivateKey.generate()
    epkc = eskc.public_key()
    nc = urandom(8)
    wallet = Wallet(ppks=self.ppks, eskc=eskc, epkc=epkc, nc=nc)

    # try:
    sock.send(wallet.epkc.public_bytes(Encoding.Raw, PublicFormat.Raw) + wallet.nc)
    data = sock.recv(40)
    # except (ConnectionError, TimeoutError):
    #     self.logger.error("failed to connect to server")
    # return False

    wallet = self.__key_exchange(data, wallet)

    self.logger.info("received keys, starting handshake")

    _, self.protocol = await asyncio.get_event_loop().create_datagram_endpoint(
        lambda: SusClientProtocol(wallet, self.protocol_id),
        sock=sock
    )
    await self.protocol.handshake_event.wait()

disconnect()

Disconnects from the server.

Source code in sus/client/client.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
def disconnect(self):
    """
    Disconnects from the server.
    """
    if not hasattr(self, "protocol"):
        self.logger.warning("not connected to server")
        return
    try:
        asyncio.get_running_loop()
        self.protocol.disconnect()
    except RuntimeError:  # not running in event loop
        pass
    self.logger.info(f"disconnected from server ({self.server_addr[0]}:{self.server_addr[1]})")

keep_alive() async

Convenience coroutine that waits until the client is disconnected.

Source code in sus/client/client.py 
156
157
158
159
160
161
162
163
164
165
166
async def keep_alive(self):
    """
    Convenience coroutine that waits until the client is disconnected.
    """
    if not hasattr(self, "protocol"):
        self.logger.warning("not connected to server")
        return
    try:
        await self.protocol.diconnection_event.wait()
    except asyncio.CancelledError:
        self.logger.info("exiting...")

send(data)

Sends a message to the server.

Parameters:

Name Type Description Default
data bytes

message to send as bytes

 required
Source code in sus/client/client.py 
132
133
134
135
136
137
138
139
140
def send(self, data: bytes):
    """
    Sends a message to the server.
    :param data: message to send as bytes
    """
    if not self.protocol:
        self.logger.warning("not connected to server")
        return
    self.protocol.send(data)

start(handlers=None) async

This coroutine is responsible for starting the client. Blocks until the client is connected. It also registers message handlers, called when a message is received.

Parameters:

Name Type Description Default
handlers Iterable[MessageHandler]
None

Returns:

Type Description
Source code in sus/client/client.py 
49
50
51
52
53
54
55
56
57
58
async def start(self, handlers: Iterable[MessageHandler] = None):
    """
    This coroutine is responsible for starting the client. Blocks until the client is connected.
    It also registers message handlers, called when a message is received.
    :param handlers:
    :return:
    """
    await self.connect()
    for handler in handlers or []:
        self.protocol.add_message_handler(handler)

Protocol

Bases: DatagramProtocol

This class is responsible for handling the UDP protocol.

Source code in sus/client/protocol.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
class SusClientProtocol(asyncio.DatagramProtocol):
    """
    This class is responsible for handling the UDP protocol.
    """
    transport: asyncio.DatagramTransport
    state: ConnectionState

    def __init__(self, wallet: Wallet, protcol_id: bytes,
                 handlers: Optional[Iterable[MessageHandler]] = None):
        """
        Initializes the client protocol.
        :param wallet: wallet containing the client's keys
        :param protcol_id: protocol ID (any bytestring)
        :param handlers: message handlers, called when a message is received
        """
        super().__init__()

        self.wallet = wallet
        self.protocol_id = protcol_id
        self.state = ConnectionState.INITIAL

        self.logger = logging.getLogger(f"sus-cl")

        self.last_seen = now()

        self.cl_enc = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + CLIENT_ENC_NONCE), None).encryptor()
        self.sr_enc = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + SERVER_ENC_NONCE), None).decryptor()
        self.cl_mac = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + CLIENT_MAC_NONCE), None).encryptor()
        self.sr_mac = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + SERVER_MAC_NONCE), None).decryptor()

        self.client_message_id = 0
        self.server_message_id = 0
        self.client_packet_id = 0
        self.server_packet_id = 0

        self.mtu_estimate = 1500

        self.message_handlers: set[MessageHandler] = set(handlers or [])
        self.handshake_event = asyncio.Event()
        self.diconnection_event = asyncio.Event()

    def add_message_handler(self, handler: MessageHandler):
        self.message_handlers.add(handler)

    def connection_made(self, transport: asyncio.DatagramTransport):
        """
        This function is called when the connection is established.
        :param transport: transport object, used to send and receive packets
        """
        self.transport = transport
        self.state = ConnectionState.HANDSHAKE
        self.send(self.protocol_id, self.wallet.token)
        self.state = ConnectionState.CONNECTED
        self.last_seen = now()
        self.logger.debug("Handshake complete")
        self.handshake_event.set()

    def datagram_received(self, data: bytes, _addr: tuple[str, int]):
        """
        This function is called when a packet is received.
        :param data: packet data
        :param _addr: originating address (always the server, unused)
        """

        match self.state:
            case ConnectionState.CONNECTED:
                pid, message = self.__verify_and_decrypt(data)
                self.logger.info(f">>> {trail_off(message.decode('utf-8')) if message else None}")
                self.handle_message(pid, message)

    def __verify_and_decrypt(self, data: bytes) -> tuple[None, None] | tuple[int, bytes]:
        """
        This function is responsible for verifying the packet and decrypting it.
        :param data: data to verify and decrypt
        :return: packet ID and message, or None if the packet is invalid
        """
        try:
            key = self.sr_mac.update(b"\x00" * 32)
            p_id = data[:8]
            payload = data[8:-16]
            tag = data[-16:]
            frame = p_id + payload
            poly1305.Poly1305.verify_tag(key, frame, tag)
        except poly1305.InvalidSignature:
            self.logger.error("Invalid signature")
            return None, None

        p_id = int.from_bytes(p_id, "little")
        message_bytes = self.sr_enc.update(payload)
        message_length = int.from_bytes(message_bytes[:4], "little")
        message: bytes = message_bytes[4:message_length + 4]
        # self.logger.info(f"Received message {p_id} ({message_length} bytes)")
        self.last_seen = now()
        self.server_packet_id = p_id
        return p_id, message

    def __split_message(self, data: bytes) -> list[bytes]:
        """
        This function is responsible for splitting a message into packets.
        :param data: data to split
        :return: list of packets
        """
        packet_length = self.mtu_estimate - 24
        return [data[i:i + packet_length] for i in range(0, len(data), packet_length)]

    def __encrypt_and_tag(self, data: bytes, token: bytes) -> list[bytes]:
        """
        This function is responsible for encrypting and tagging a message. Uses the ChaCha20-Poly1305 AEAD to
        encrypt and authenticate the message.
        :param data: data to encrypt
        :param token: optional token to include in the first packet
        :return: packets containing the encrypted and tagged message to send to the server
        """
        message_bytes = len(data).to_bytes(4, "little") + data
        packet_length = self.mtu_estimate - 24
        padded_message_bytes = message_bytes + b"\x00" * (
                packet_length - ((len(message_bytes) + len(token)) % packet_length))

        ciphertext = self.cl_enc.update(padded_message_bytes)
        self.logger.debug(f"--- {trail_off(ciphertext.hex())}")
        if token:
            self.logger.debug(f"TOK {trail_off(token.hex())}")

        payloads = self.__split_message(token + ciphertext)

        packets = []
        for payload in payloads:
            key = self.cl_mac.update(b"\x00" * 32)
            p_id = self.client_packet_id.to_bytes(8, "little")
            frame = p_id + payload
            tag = poly1305.Poly1305.generate_tag(key, frame)
            packets.append(frame + tag)
            self.client_packet_id += 1
        return packets

    def send(self, data: bytes, token: bytes = b""):
        """
        This function is responsible for sending a message to the server.
        :param data: data to send
        :param token: token to include in the first packet. DO NOT INCLUDE THE TOKEN IN SUBSEQUENT PACKETS --
                      a MalformedPacket error will be raised!
        :raises MalformedPacket: if the token is included in a subsequent packet
        """
        if self.state not in (ConnectionState.CONNECTED, ConnectionState.HANDSHAKE):
            return
        if token and self.client_packet_id != 0:
            raise MalformedPacket("Token can only be included in the first packet")
        self.logger.info(f"<<< {trail_off(data.decode('utf-8'))}")
        packets = self.__encrypt_and_tag(data, token)
        self.logger.info(f"Sending {len(data)} bytes in {len(packets)} packets")
        for packet in packets:
            self.transport.sendto(packet, None)

    def disconnect(self):
        """
        Disconnects the client from the server.
        """
        self.logger.warning("Disconnecting from server...")
        self.transport.close()

    def connection_lost(self, exc):
        """
        Called when the connection is lost. Sets the disconnection event.
        :param exc: exception raised, if any
        """
        self.logger.warning("Connection to server lost")
        if exc:
            self.logger.exception(exc)
            self.state = ConnectionState.ERROR
        else:
            self.state = ConnectionState.DISCONNECTED
        self.diconnection_event.set()

    def handle_message(self, pid: int, message: bytes):
        """
        Calls all message handlers asynchronously.
        :param pid: packet ID
        :param message: message bytes
        :return:
        """
        asyncio.gather(*[handler(("", 0), pid, message) for handler in self.message_handlers])

__encrypt_and_tag(data, token)

This function is responsible for encrypting and tagging a message. Uses the ChaCha20-Poly1305 AEAD to encrypt and authenticate the message.

Parameters:

Name Type Description Default
data bytes

data to encrypt

 required
token bytes

optional token to include in the first packet

 required

Returns:

Type Description
list[bytes]

packets containing the encrypted and tagged message to send to the server
Source code in sus/client/protocol.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def __encrypt_and_tag(self, data: bytes, token: bytes) -> list[bytes]:
    """
    This function is responsible for encrypting and tagging a message. Uses the ChaCha20-Poly1305 AEAD to
    encrypt and authenticate the message.
    :param data: data to encrypt
    :param token: optional token to include in the first packet
    :return: packets containing the encrypted and tagged message to send to the server
    """
    message_bytes = len(data).to_bytes(4, "little") + data
    packet_length = self.mtu_estimate - 24
    padded_message_bytes = message_bytes + b"\x00" * (
            packet_length - ((len(message_bytes) + len(token)) % packet_length))

    ciphertext = self.cl_enc.update(padded_message_bytes)
    self.logger.debug(f"--- {trail_off(ciphertext.hex())}")
    if token:
        self.logger.debug(f"TOK {trail_off(token.hex())}")

    payloads = self.__split_message(token + ciphertext)

    packets = []
    for payload in payloads:
        key = self.cl_mac.update(b"\x00" * 32)
        p_id = self.client_packet_id.to_bytes(8, "little")
        frame = p_id + payload
        tag = poly1305.Poly1305.generate_tag(key, frame)
        packets.append(frame + tag)
        self.client_packet_id += 1
    return packets

__init__(wallet, protcol_id, handlers=None)

Initializes the client protocol.

Parameters:

Name Type Description Default
wallet Wallet

wallet containing the client's keys

 required
protcol_id bytes

protocol ID (any bytestring)

 required
handlers Optional[Iterable[MessageHandler]]

message handlers, called when a message is received

 None
Source code in sus/client/protocol.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(self, wallet: Wallet, protcol_id: bytes,
             handlers: Optional[Iterable[MessageHandler]] = None):
    """
    Initializes the client protocol.
    :param wallet: wallet containing the client's keys
    :param protcol_id: protocol ID (any bytestring)
    :param handlers: message handlers, called when a message is received
    """
    super().__init__()

    self.wallet = wallet
    self.protocol_id = protcol_id
    self.state = ConnectionState.INITIAL

    self.logger = logging.getLogger(f"sus-cl")

    self.last_seen = now()

    self.cl_enc = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + CLIENT_ENC_NONCE), None).encryptor()
    self.sr_enc = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + SERVER_ENC_NONCE), None).decryptor()
    self.cl_mac = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + CLIENT_MAC_NONCE), None).encryptor()
    self.sr_mac = Cipher(ChaCha20(wallet.shared_secret, b"\x00" * 8 + SERVER_MAC_NONCE), None).decryptor()

    self.client_message_id = 0
    self.server_message_id = 0
    self.client_packet_id = 0
    self.server_packet_id = 0

    self.mtu_estimate = 1500

    self.message_handlers: set[MessageHandler] = set(handlers or [])
    self.handshake_event = asyncio.Event()
    self.diconnection_event = asyncio.Event()

__split_message(data)

This function is responsible for splitting a message into packets.

Parameters:

Name Type Description Default
data bytes

data to split

 required

Returns:

Type Description
list[bytes]

list of packets
Source code in sus/client/protocol.py 
116
117
118
119
120
121
122
123
def __split_message(self, data: bytes) -> list[bytes]:
    """
    This function is responsible for splitting a message into packets.
    :param data: data to split
    :return: list of packets
    """
    packet_length = self.mtu_estimate - 24
    return [data[i:i + packet_length] for i in range(0, len(data), packet_length)]

__verify_and_decrypt(data)

This function is responsible for verifying the packet and decrypting it.

Parameters:

Name Type Description Default
data bytes

data to verify and decrypt

 required

Returns:

Type Description
tuple[None, None] | tuple[int, bytes]

packet ID and message, or None if the packet is invalid
Source code in sus/client/protocol.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def __verify_and_decrypt(self, data: bytes) -> tuple[None, None] | tuple[int, bytes]:
    """
    This function is responsible for verifying the packet and decrypting it.
    :param data: data to verify and decrypt
    :return: packet ID and message, or None if the packet is invalid
    """
    try:
        key = self.sr_mac.update(b"\x00" * 32)
        p_id = data[:8]
        payload = data[8:-16]
        tag = data[-16:]
        frame = p_id + payload
        poly1305.Poly1305.verify_tag(key, frame, tag)
    except poly1305.InvalidSignature:
        self.logger.error("Invalid signature")
        return None, None

    p_id = int.from_bytes(p_id, "little")
    message_bytes = self.sr_enc.update(payload)
    message_length = int.from_bytes(message_bytes[:4], "little")
    message: bytes = message_bytes[4:message_length + 4]
    # self.logger.info(f"Received message {p_id} ({message_length} bytes)")
    self.last_seen = now()
    self.server_packet_id = p_id
    return p_id, message

connection_lost(exc)

Called when the connection is lost. Sets the disconnection event.

Parameters:

Name Type Description Default
exc

exception raised, if any

 required
Source code in sus/client/protocol.py 
180
181
182
183
184
185
186
187
188
189
190
191
def connection_lost(self, exc):
    """
    Called when the connection is lost. Sets the disconnection event.
    :param exc: exception raised, if any
    """
    self.logger.warning("Connection to server lost")
    if exc:
        self.logger.exception(exc)
        self.state = ConnectionState.ERROR
    else:
        self.state = ConnectionState.DISCONNECTED
    self.diconnection_event.set()

connection_made(transport)

This function is called when the connection is established.

Parameters:

Name Type Description Default
transport DatagramTransport

transport object, used to send and receive packets

 required
Source code in sus/client/protocol.py 
64
65
66
67
68
69
70
71
72
73
74
75
def connection_made(self, transport: asyncio.DatagramTransport):
    """
    This function is called when the connection is established.
    :param transport: transport object, used to send and receive packets
    """
    self.transport = transport
    self.state = ConnectionState.HANDSHAKE
    self.send(self.protocol_id, self.wallet.token)
    self.state = ConnectionState.CONNECTED
    self.last_seen = now()
    self.logger.debug("Handshake complete")
    self.handshake_event.set()

datagram_received(data, _addr)

This function is called when a packet is received.

Parameters:

Name Type Description Default
data bytes

packet data

 required
_addr tuple[str, int]

originating address (always the server, unused)

 required
Source code in sus/client/protocol.py 
77
78
79
80
81
82
83
84
85
86
87
88
def datagram_received(self, data: bytes, _addr: tuple[str, int]):
    """
    This function is called when a packet is received.
    :param data: packet data
    :param _addr: originating address (always the server, unused)
    """

    match self.state:
        case ConnectionState.CONNECTED:
            pid, message = self.__verify_and_decrypt(data)
            self.logger.info(f">>> {trail_off(message.decode('utf-8')) if message else None}")
            self.handle_message(pid, message)

disconnect()

Disconnects the client from the server.

Source code in sus/client/protocol.py 
173
174
175
176
177
178
def disconnect(self):
    """
    Disconnects the client from the server.
    """
    self.logger.warning("Disconnecting from server...")
    self.transport.close()

handle_message(pid, message)

Calls all message handlers asynchronously.

Parameters:

Name Type Description Default
pid int

packet ID

 required
message bytes

message bytes

 required

Returns:

Type Description
Source code in sus/client/protocol.py 
193
194
195
196
197
198
199
200
def handle_message(self, pid: int, message: bytes):
    """
    Calls all message handlers asynchronously.
    :param pid: packet ID
    :param message: message bytes
    :return:
    """
    asyncio.gather(*[handler(("", 0), pid, message) for handler in self.message_handlers])

send(data, token=b'')

This function is responsible for sending a message to the server.

Parameters:

Name Type Description Default
data bytes

data to send

 required
token bytes

token to include in the first packet. DO NOT INCLUDE THE TOKEN IN SUBSEQUENT PACKETS -- a MalformedPacket error will be raised!

 b''

Raises:

Type Description
MalformedPacket

if the token is included in a subsequent packet
Source code in sus/client/protocol.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def send(self, data: bytes, token: bytes = b""):
    """
    This function is responsible for sending a message to the server.
    :param data: data to send
    :param token: token to include in the first packet. DO NOT INCLUDE THE TOKEN IN SUBSEQUENT PACKETS --
                  a MalformedPacket error will be raised!
    :raises MalformedPacket: if the token is included in a subsequent packet
    """
    if self.state not in (ConnectionState.CONNECTED, ConnectionState.HANDSHAKE):
        return
    if token and self.client_packet_id != 0:
        raise MalformedPacket("Token can only be included in the first packet")
    self.logger.info(f"<<< {trail_off(data.decode('utf-8'))}")
    packets = self.__encrypt_and_tag(data, token)
    self.logger.info(f"Sending {len(data)} bytes in {len(packets)} packets")
    for packet in packets:
        self.transport.sendto(packet, None)