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# Crypto support
17#
18# See Bluetooth spec Vol 3, Part H - 2.2 CRYPTOGRAPHIC TOOLBOX
19# -----------------------------------------------------------------------------
20
21# -----------------------------------------------------------------------------
22# Imports
23# -----------------------------------------------------------------------------
24from __future__ import annotations
25
26import logging
27import operator
28
29import secrets
30from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
31from cryptography.hazmat.primitives.asymmetric.ec import (
32    generate_private_key,
33    ECDH,
34    EllipticCurvePrivateKey,
35    EllipticCurvePublicNumbers,
36    EllipticCurvePrivateNumbers,
37    SECP256R1,
38)
39from cryptography.hazmat.primitives import cmac
40from typing import Tuple
41
42
43# -----------------------------------------------------------------------------
44# Logging
45# -----------------------------------------------------------------------------
46logger = logging.getLogger(__name__)
47
48
49# -----------------------------------------------------------------------------
50# Classes
51# -----------------------------------------------------------------------------
52class EccKey:
53    def __init__(self, private_key: EllipticCurvePrivateKey) -> None:
54        self.private_key = private_key
55
56    @classmethod
57    def generate(cls) -> EccKey:
58        private_key = generate_private_key(SECP256R1())
59        return cls(private_key)
60
61    @classmethod
62    def from_private_key_bytes(
63        cls, d_bytes: bytes, x_bytes: bytes, y_bytes: bytes
64    ) -> EccKey:
65        d = int.from_bytes(d_bytes, byteorder='big', signed=False)
66        x = int.from_bytes(x_bytes, byteorder='big', signed=False)
67        y = int.from_bytes(y_bytes, byteorder='big', signed=False)
68        private_key = EllipticCurvePrivateNumbers(
69            d, EllipticCurvePublicNumbers(x, y, SECP256R1())
70        ).private_key()
71        return cls(private_key)
72
73    @property
74    def x(self) -> bytes:
75        return (
76            self.private_key.public_key()
77            .public_numbers()
78            .x.to_bytes(32, byteorder='big')
79        )
80
81    @property
82    def y(self) -> bytes:
83        return (
84            self.private_key.public_key()
85            .public_numbers()
86            .y.to_bytes(32, byteorder='big')
87        )
88
89    def dh(self, public_key_x: bytes, public_key_y: bytes) -> bytes:
90        x = int.from_bytes(public_key_x, byteorder='big', signed=False)
91        y = int.from_bytes(public_key_y, byteorder='big', signed=False)
92        public_key = EllipticCurvePublicNumbers(x, y, SECP256R1()).public_key()
93        shared_key = self.private_key.exchange(ECDH(), public_key)
94
95        return shared_key
96
97
98# -----------------------------------------------------------------------------
99# Functions
100# -----------------------------------------------------------------------------
101
102
103# -----------------------------------------------------------------------------
104def generate_prand() -> bytes:
105    '''Generates random 3 bytes, with the 2 most significant bits of 0b01.
106
107    See Bluetooth spec, Vol 6, Part E - Table 1.2.
108    '''
109    prand_bytes = secrets.token_bytes(6)
110    return prand_bytes[:2] + bytes([(prand_bytes[2] & 0b01111111) | 0b01000000])
111
112
113# -----------------------------------------------------------------------------
114def xor(x: bytes, y: bytes) -> bytes:
115    assert len(x) == len(y)
116    return bytes(map(operator.xor, x, y))
117
118
119# -----------------------------------------------------------------------------
120def reverse(input: bytes) -> bytes:
121    '''
122    Returns bytes of input in reversed endianness.
123    '''
124    return input[::-1]
125
126
127# -----------------------------------------------------------------------------
128def r() -> bytes:
129    '''
130    Generate 16 bytes of random data
131    '''
132    return secrets.token_bytes(16)
133
134
135# -----------------------------------------------------------------------------
136def e(key: bytes, data: bytes) -> bytes:
137    '''
138    AES-128 ECB, expecting byte-swapped inputs and producing a byte-swapped output.
139
140    See Bluetooth spec Vol 3, Part H - 2.2.1 Security function e
141    '''
142
143    cipher = Cipher(algorithms.AES(reverse(key)), modes.ECB())
144    encryptor = cipher.encryptor()
145    return reverse(encryptor.update(reverse(data)))
146
147
148# -----------------------------------------------------------------------------
149def ah(k: bytes, r: bytes) -> bytes:  # pylint: disable=redefined-outer-name
150    '''
151    See Bluetooth spec Vol 3, Part H - 2.2.2 Random Address Hash function ah
152    '''
153
154    padding = bytes(13)
155    r_prime = r + padding
156    return e(k, r_prime)[0:3]
157
158
159# -----------------------------------------------------------------------------
160def c1(
161    k: bytes,
162    r: bytes,
163    preq: bytes,
164    pres: bytes,
165    iat: int,
166    rat: int,
167    ia: bytes,
168    ra: bytes,
169) -> bytes:  # pylint: disable=redefined-outer-name
170    '''
171    See Bluetooth spec, Vol 3, Part H - 2.2.3 Confirm value generation function c1 for
172    LE Legacy Pairing
173    '''
174
175    p1 = bytes([iat, rat]) + preq + pres
176    p2 = ra + ia + bytes([0, 0, 0, 0])
177    return e(k, xor(e(k, xor(r, p1)), p2))
178
179
180# -----------------------------------------------------------------------------
181def s1(k: bytes, r1: bytes, r2: bytes) -> bytes:
182    '''
183    See Bluetooth spec, Vol 3, Part H - 2.2.4 Key generation function s1 for LE Legacy
184    Pairing
185    '''
186
187    return e(k, r2[0:8] + r1[0:8])
188
189
190# -----------------------------------------------------------------------------
191def aes_cmac(m: bytes, k: bytes) -> bytes:
192    '''
193    See Bluetooth spec, Vol 3, Part H - 2.2.5 FunctionAES-CMAC
194
195    NOTE: the input and output of this internal function are in big-endian byte order
196    '''
197    mac = cmac.CMAC(algorithms.AES(k))
198    mac.update(m)
199    return mac.finalize()
200
201
202# -----------------------------------------------------------------------------
203def f4(u: bytes, v: bytes, x: bytes, z: bytes) -> bytes:
204    '''
205    See Bluetooth spec, Vol 3, Part H - 2.2.6 LE Secure Connections Confirm Value
206    Generation Function f4
207    '''
208    return reverse(aes_cmac(reverse(u) + reverse(v) + z, reverse(x)))
209
210
211# -----------------------------------------------------------------------------
212def f5(w: bytes, n1: bytes, n2: bytes, a1: bytes, a2: bytes) -> Tuple[bytes, bytes]:
213    '''
214    See Bluetooth spec, Vol 3, Part H - 2.2.7 LE Secure Connections Key Generation
215    Function f5
216
217    NOTE: this returns a tuple: (MacKey, LTK) in little-endian byte order
218    '''
219    salt = bytes.fromhex('6C888391AAF5A53860370BDB5A6083BE')
220    t = aes_cmac(reverse(w), salt)
221    key_id = bytes([0x62, 0x74, 0x6C, 0x65])
222    return (
223        reverse(
224            aes_cmac(
225                bytes([0])
226                + key_id
227                + reverse(n1)
228                + reverse(n2)
229                + reverse(a1)
230                + reverse(a2)
231                + bytes([1, 0]),
232                t,
233            )
234        ),
235        reverse(
236            aes_cmac(
237                bytes([1])
238                + key_id
239                + reverse(n1)
240                + reverse(n2)
241                + reverse(a1)
242                + reverse(a2)
243                + bytes([1, 0]),
244                t,
245            )
246        ),
247    )
248
249
250# -----------------------------------------------------------------------------
251def f6(
252    w: bytes, n1: bytes, n2: bytes, r: bytes, io_cap: bytes, a1: bytes, a2: bytes
253) -> bytes:  # pylint: disable=redefined-outer-name
254    '''
255    See Bluetooth spec, Vol 3, Part H - 2.2.8 LE Secure Connections Check Value
256    Generation Function f6
257    '''
258    return reverse(
259        aes_cmac(
260            reverse(n1)
261            + reverse(n2)
262            + reverse(r)
263            + reverse(io_cap)
264            + reverse(a1)
265            + reverse(a2),
266            reverse(w),
267        )
268    )
269
270
271# -----------------------------------------------------------------------------
272def g2(u: bytes, v: bytes, x: bytes, y: bytes) -> int:
273    '''
274    See Bluetooth spec, Vol 3, Part H - 2.2.9 LE Secure Connections Numeric Comparison
275    Value Generation Function g2
276    '''
277    return int.from_bytes(
278        aes_cmac(
279            reverse(u) + reverse(v) + reverse(y),
280            reverse(x),
281        )[-4:],
282        byteorder='big',
283    )
284
285
286# -----------------------------------------------------------------------------
287def h6(w: bytes, key_id: bytes) -> bytes:
288    '''
289    See Bluetooth spec, Vol 3, Part H - 2.2.10 Link key conversion function h6
290    '''
291    return reverse(aes_cmac(key_id, reverse(w)))
292
293
294# -----------------------------------------------------------------------------
295def h7(salt: bytes, w: bytes) -> bytes:
296    '''
297    See Bluetooth spec, Vol 3, Part H - 2.2.11 Link key conversion function h7
298    '''
299    return reverse(aes_cmac(reverse(w), salt))
300