1# Copyright 2021-2022 Google LLC
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#      https://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
15# -----------------------------------------------------------------------------
16# Imports
17# -----------------------------------------------------------------------------
18import asyncio
19import logging
20import os
21import click
22
23from bumble import l2cap
24from bumble.colors import color
25from bumble.transport import open_transport_or_link
26from bumble.device import Device
27from bumble.utils import FlowControlAsyncPipe
28from bumble.hci import HCI_Constant
29
30
31# -----------------------------------------------------------------------------
32class ServerBridge:
33    """
34    L2CAP CoC server bridge: waits for a peer to connect an L2CAP CoC channel
35    on a specified PSM. When the connection is made, the bridge connects a TCP
36    socket to a remote host and bridges the data in both directions, with flow
37    control.
38    When the L2CAP CoC channel is closed, the bridge disconnects the TCP socket
39    and waits for a new L2CAP CoC channel to be connected.
40    When the TCP connection is closed by the TCP server, XXXX
41    """
42
43    def __init__(self, psm, max_credits, mtu, mps, tcp_host, tcp_port):
44        self.psm = psm
45        self.max_credits = max_credits
46        self.mtu = mtu
47        self.mps = mps
48        self.tcp_host = tcp_host
49        self.tcp_port = tcp_port
50
51    async def start(self, device: Device) -> None:
52        # Listen for incoming L2CAP channel connections
53        device.create_l2cap_server(
54            spec=l2cap.LeCreditBasedChannelSpec(
55                psm=self.psm, mtu=self.mtu, mps=self.mps, max_credits=self.max_credits
56            ),
57            handler=self.on_channel,
58        )
59        print(
60            color(f'### Listening for channel connection on PSM {self.psm}', 'yellow')
61        )
62
63        def on_ble_connection(connection):
64            def on_ble_disconnection(reason):
65                print(
66                    color('@@@ Bluetooth disconnection:', 'red'),
67                    HCI_Constant.error_name(reason),
68                )
69
70            print(color('@@@ Bluetooth connection:', 'green'), connection)
71            connection.on('disconnection', on_ble_disconnection)
72
73        device.on('connection', on_ble_connection)
74
75        await device.start_advertising(auto_restart=True)
76
77    # Called when a new L2CAP connection is established
78    def on_channel(self, l2cap_channel):
79        print(color('*** L2CAP channel:', 'cyan'), l2cap_channel)
80
81        class Pipe:
82            def __init__(self, bridge, l2cap_channel):
83                self.bridge = bridge
84                self.tcp_transport = None
85                self.l2cap_channel = l2cap_channel
86
87                l2cap_channel.on('close', self.on_l2cap_close)
88                l2cap_channel.sink = self.on_channel_sdu
89
90            async def connect_to_tcp(self):
91                # Connect to the TCP server
92                print(
93                    color(
94                        f'### Connecting to TCP {self.bridge.tcp_host}:'
95                        f'{self.bridge.tcp_port}...',
96                        'yellow',
97                    )
98                )
99
100                class TcpClientProtocol(asyncio.Protocol):
101                    def __init__(self, pipe):
102                        self.pipe = pipe
103
104                    def connection_lost(self, exc):
105                        print(color(f'!!! TCP connection lost: {exc}', 'red'))
106                        if self.pipe.l2cap_channel is not None:
107                            asyncio.create_task(self.pipe.l2cap_channel.disconnect())
108
109                    def data_received(self, data):
110                        print(color(f'<<< [TCP DATA]: {len(data)} bytes', 'blue'))
111                        self.pipe.l2cap_channel.write(data)
112
113                try:
114                    (
115                        self.tcp_transport,
116                        _,
117                    ) = await asyncio.get_running_loop().create_connection(
118                        lambda: TcpClientProtocol(self),
119                        host=self.bridge.tcp_host,
120                        port=self.bridge.tcp_port,
121                    )
122                    print(color('### Connected', 'green'))
123                except Exception as error:
124                    print(color(f'!!! Connection failed: {error}', 'red'))
125                    await self.l2cap_channel.disconnect()
126
127            def on_l2cap_close(self):
128                print(color('*** L2CAP channel closed', 'red'))
129                self.l2cap_channel = None
130                if self.tcp_transport is not None:
131                    self.tcp_transport.close()
132
133            def on_channel_sdu(self, sdu):
134                print(color(f'<<< [L2CAP SDU]: {len(sdu)} bytes', 'cyan'))
135                if self.tcp_transport is None:
136                    print(color('!!! TCP socket not open, dropping', 'red'))
137                    return
138                self.tcp_transport.write(sdu)
139
140        pipe = Pipe(self, l2cap_channel)
141
142        asyncio.create_task(pipe.connect_to_tcp())
143
144
145# -----------------------------------------------------------------------------
146class ClientBridge:
147    """
148    L2CAP CoC client bridge: connects to a BLE device, then waits for an inbound
149    TCP connection on a specified port number. When a TCP client connects, an
150    L2CAP CoC channel connection to the BLE device is established, and the data
151    is bridged in both directions, with flow control.
152    When the TCP connection is closed by the client, the L2CAP CoC channel is
153    disconnected, but the connection to the BLE device remains, ready for a new
154    TCP client to connect.
155    When the L2CAP CoC channel is closed, XXXX
156    """
157
158    READ_CHUNK_SIZE = 4096
159
160    def __init__(self, psm, max_credits, mtu, mps, address, tcp_host, tcp_port):
161        self.psm = psm
162        self.max_credits = max_credits
163        self.mtu = mtu
164        self.mps = mps
165        self.address = address
166        self.tcp_host = tcp_host
167        self.tcp_port = tcp_port
168
169    async def start(self, device):
170        print(color(f'### Connecting to {self.address}...', 'yellow'))
171        connection = await device.connect(self.address)
172        print(color('### Connected', 'green'))
173
174        # Called when the BLE connection is disconnected
175        def on_ble_disconnection(reason):
176            print(
177                color('@@@ Bluetooth disconnection:', 'red'),
178                HCI_Constant.error_name(reason),
179            )
180
181        connection.on('disconnection', on_ble_disconnection)
182
183        # Called when a TCP connection is established
184        async def on_tcp_connection(reader, writer):
185            peer_name = writer.get_extra_info('peer_name')
186            print(color(f'<<< TCP connection from {peer_name}', 'magenta'))
187
188            def on_channel_sdu(sdu):
189                print(color(f'<<< [L2CAP SDU]: {len(sdu)} bytes', 'cyan'))
190                l2cap_to_tcp_pipe.write(sdu)
191
192            def on_l2cap_close():
193                print(color('*** L2CAP channel closed', 'red'))
194                l2cap_to_tcp_pipe.stop()
195                writer.close()
196
197            # Connect a new L2CAP channel
198            print(color(f'>>> Opening L2CAP channel on PSM = {self.psm}', 'yellow'))
199            try:
200                l2cap_channel = await connection.create_l2cap_channel(
201                    spec=l2cap.LeCreditBasedChannelSpec(
202                        psm=self.psm,
203                        max_credits=self.max_credits,
204                        mtu=self.mtu,
205                        mps=self.mps,
206                    )
207                )
208                print(color('*** L2CAP channel:', 'cyan'), l2cap_channel)
209            except Exception as error:
210                print(color(f'!!! Connection failed: {error}', 'red'))
211                writer.close()
212                return
213
214            l2cap_channel.sink = on_channel_sdu
215            l2cap_channel.on('close', on_l2cap_close)
216
217            # Start a flow control pipe from L2CAP to TCP
218            l2cap_to_tcp_pipe = FlowControlAsyncPipe(
219                l2cap_channel.pause_reading,
220                l2cap_channel.resume_reading,
221                writer.write,
222                writer.drain,
223            )
224            l2cap_to_tcp_pipe.start()
225
226            # Pipe data from TCP to L2CAP
227            while True:
228                try:
229                    data = await reader.read(self.READ_CHUNK_SIZE)
230
231                    if len(data) == 0:
232                        print(color('!!! End of stream', 'red'))
233                        await l2cap_channel.disconnect()
234                        return
235
236                    print(color(f'<<< [TCP DATA]: {len(data)} bytes', 'blue'))
237                    l2cap_channel.write(data)
238                    await l2cap_channel.drain()
239                except Exception as error:
240                    print(f'!!! Exception: {error}')
241                    break
242
243            writer.close()
244            print(color('~~~ Bye bye', 'magenta'))
245
246        await asyncio.start_server(
247            on_tcp_connection,
248            host=self.tcp_host if self.tcp_host != '_' else None,
249            port=self.tcp_port,
250        )
251        print(
252            color(
253                f'### Listening for TCP connections on port {self.tcp_port}', 'magenta'
254            )
255        )
256
257
258# -----------------------------------------------------------------------------
259async def run(device_config, hci_transport, bridge):
260    print('<<< connecting to HCI...')
261    async with await open_transport_or_link(hci_transport) as (hci_source, hci_sink):
262        print('<<< connected')
263
264        device = Device.from_config_file_with_hci(device_config, hci_source, hci_sink)
265
266        # Let's go
267        await device.power_on()
268        await bridge.start(device)
269
270        # Wait until the transport terminates
271        await hci_source.wait_for_termination()
272
273
274# -----------------------------------------------------------------------------
275@click.group()
276@click.pass_context
277@click.option('--device-config', help='Device configuration file', required=True)
278@click.option('--hci-transport', help='HCI transport', required=True)
279@click.option('--psm', help='PSM for L2CAP', type=int, default=1234)
280@click.option(
281    '--l2cap-max-credits',
282    help='Maximum L2CAP Credits',
283    type=click.IntRange(1, 65535),
284    default=128,
285)
286@click.option(
287    '--l2cap-mtu',
288    help='L2CAP MTU',
289    type=click.IntRange(
290        l2cap.L2CAP_LE_CREDIT_BASED_CONNECTION_MIN_MTU,
291        l2cap.L2CAP_LE_CREDIT_BASED_CONNECTION_MAX_MTU,
292    ),
293    default=1024,
294)
295@click.option(
296    '--l2cap-mps',
297    help='L2CAP MPS',
298    type=click.IntRange(
299        l2cap.L2CAP_LE_CREDIT_BASED_CONNECTION_MIN_MPS,
300        l2cap.L2CAP_LE_CREDIT_BASED_CONNECTION_MAX_MPS,
301    ),
302    default=1024,
303)
304def cli(
305    context,
306    device_config,
307    hci_transport,
308    psm,
309    l2cap_max_credits,
310    l2cap_mtu,
311    l2cap_mps,
312):
313    context.ensure_object(dict)
314    context.obj['device_config'] = device_config
315    context.obj['hci_transport'] = hci_transport
316    context.obj['psm'] = psm
317    context.obj['max_credits'] = l2cap_max_credits
318    context.obj['mtu'] = l2cap_mtu
319    context.obj['mps'] = l2cap_mps
320
321
322# -----------------------------------------------------------------------------
323@cli.command()
324@click.pass_context
325@click.option('--tcp-host', help='TCP host', default='localhost')
326@click.option('--tcp-port', help='TCP port', default=9544)
327def server(context, tcp_host, tcp_port):
328    bridge = ServerBridge(
329        context.obj['psm'],
330        context.obj['max_credits'],
331        context.obj['mtu'],
332        context.obj['mps'],
333        tcp_host,
334        tcp_port,
335    )
336    asyncio.run(run(context.obj['device_config'], context.obj['hci_transport'], bridge))
337
338
339# -----------------------------------------------------------------------------
340@cli.command()
341@click.pass_context
342@click.argument('bluetooth-address')
343@click.option('--tcp-host', help='TCP host', default='_')
344@click.option('--tcp-port', help='TCP port', default=9543)
345def client(context, bluetooth_address, tcp_host, tcp_port):
346    bridge = ClientBridge(
347        context.obj['psm'],
348        context.obj['max_credits'],
349        context.obj['mtu'],
350        context.obj['mps'],
351        bluetooth_address,
352        tcp_host,
353        tcp_port,
354    )
355    asyncio.run(run(context.obj['device_config'], context.obj['hci_transport'], bridge))
356
357
358# -----------------------------------------------------------------------------
359logging.basicConfig(level=os.environ.get('BUMBLE_LOGLEVEL', 'WARNING').upper())
360if __name__ == '__main__':
361    cli(obj={})  # pylint: disable=no-value-for-parameter
362