codekingpro/portable-devtools
114k
1#define PY_SSIZE_T_CLEAN
2
3#include <Python.h>
4#include <openssl/err.h>
5#include <openssl/evp.h>
6
7#define MODULE_NAME "aioquic._crypto"
8
9#define AEAD_KEY_LENGTH_MAX 32
10#define AEAD_NONCE_LENGTH 12
11#define AEAD_TAG_LENGTH 16
12
13#define PACKET_LENGTH_MAX 1500
14#define PACKET_NUMBER_LENGTH_MAX 4
15#define SAMPLE_LENGTH 16
16
17#define CHECK_RESULT(expr) \
18 if (!(expr)) { \
19 ERR_clear_error(); \
20 PyErr_SetString(CryptoError, "OpenSSL call failed"); \
21 return NULL; \
22 }
23
24#define CHECK_RESULT_CTOR(expr) \
25 if (!(expr)) { \
26 ERR_clear_error(); \
27 PyErr_SetString(CryptoError, "OpenSSL call failed"); \
28 return -1; \
29 }
30
31static PyObject *CryptoError;
32
33/* AEAD */
34
35typedef struct {
36 PyObject_HEAD
37 EVP_CIPHER_CTX *decrypt_ctx;
38 EVP_CIPHER_CTX *encrypt_ctx;
39 unsigned char buffer[PACKET_LENGTH_MAX];
40 unsigned char key[AEAD_KEY_LENGTH_MAX];
41 unsigned char iv[AEAD_NONCE_LENGTH];
42 unsigned char nonce[AEAD_NONCE_LENGTH];
43} AEADObject;
44
45static PyObject *AEADType;
46
47static EVP_CIPHER_CTX *
48create_ctx(const EVP_CIPHER *cipher, int key_length, int operation)
49{
50 EVP_CIPHER_CTX *ctx;
51 int res;
52
53 ctx = EVP_CIPHER_CTX_new();
54 CHECK_RESULT(ctx != 0);
55
56 res = EVP_CipherInit_ex(ctx, cipher, NULL, NULL, NULL, operation);
57 CHECK_RESULT(res != 0);
58
59 res = EVP_CIPHER_CTX_set_key_length(ctx, key_length);
60 CHECK_RESULT(res != 0);
61
62 res = EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_CCM_SET_IVLEN, AEAD_NONCE_LENGTH, NULL);
63 CHECK_RESULT(res != 0);
64
65 return ctx;
66}
67
68static int
69AEAD_init(AEADObject *self, PyObject *args, PyObject *kwargs)
70{
71 const char *cipher_name;
72 const unsigned char *key, *iv;
73 Py_ssize_t cipher_name_len, key_len, iv_len;
74
75 if (!PyArg_ParseTuple(args, "y#y#y#", &cipher_name, &cipher_name_len, &key, &key_len, &iv, &iv_len))
76 return -1;
77
78 const EVP_CIPHER *evp_cipher = EVP_get_cipherbyname(cipher_name);
79 if (evp_cipher == 0) {
80 PyErr_Format(CryptoError, "Invalid cipher name: %s", cipher_name);
81 return -1;
82 }
83 if (key_len > AEAD_KEY_LENGTH_MAX) {
84 PyErr_SetString(CryptoError, "Invalid key length");
85 return -1;
86 }
87 if (iv_len > AEAD_NONCE_LENGTH) {
88 PyErr_SetString(CryptoError, "Invalid iv length");
89 return -1;
90 }
91
92 memcpy(self->key, key, key_len);
93 memcpy(self->iv, iv, iv_len);
94
95 self->decrypt_ctx = create_ctx(evp_cipher, key_len, 0);
96 CHECK_RESULT_CTOR(self->decrypt_ctx != 0);
97
98 self->encrypt_ctx = create_ctx(evp_cipher, key_len, 1);
99 CHECK_RESULT_CTOR(self->encrypt_ctx != 0);
100
101 return 0;
102}
103
104static void
105AEAD_dealloc(AEADObject *self)
106{
107 EVP_CIPHER_CTX_free(self->decrypt_ctx);
108 EVP_CIPHER_CTX_free(self->encrypt_ctx);
109 PyTypeObject *tp = Py_TYPE(self);
110 freefunc free = PyType_GetSlot(tp, Py_tp_free);
111 free(self);
112 Py_DECREF(tp);
113}
114
115static PyObject*
116AEAD_decrypt(AEADObject *self, PyObject *args)
117{
118 const unsigned char *data, *associated;
119 Py_ssize_t data_len, associated_len;
120 int outlen, outlen2, res;
121 uint64_t pn;
122
123 if (!PyArg_ParseTuple(args, "y#y#K", &data, &data_len, &associated, &associated_len, &pn))
124 return NULL;
125
126 if (data_len < AEAD_TAG_LENGTH || data_len > PACKET_LENGTH_MAX) {
127 PyErr_SetString(CryptoError, "Invalid payload length");
128 return NULL;
129 }
130
131 memcpy(self->nonce, self->iv, AEAD_NONCE_LENGTH);
132 for (int i = 0; i < 8; ++i) {
133 self->nonce[AEAD_NONCE_LENGTH - 1 - i] ^= (uint8_t)(pn >> 8 * i);
134 }
135
136 res = EVP_CIPHER_CTX_ctrl(self->decrypt_ctx, EVP_CTRL_CCM_SET_TAG, AEAD_TAG_LENGTH, (void*)(data + (data_len - AEAD_TAG_LENGTH)));
137 CHECK_RESULT(res != 0);
138
139 res = EVP_CipherInit_ex(self->decrypt_ctx, NULL, NULL, self->key, self->nonce, 0);
140 CHECK_RESULT(res != 0);
141
142 res = EVP_CipherUpdate(self->decrypt_ctx, NULL, &outlen, associated, associated_len);
143 CHECK_RESULT(res != 0);
144
145 res = EVP_CipherUpdate(self->decrypt_ctx, self->buffer, &outlen, data, data_len - AEAD_TAG_LENGTH);
146 CHECK_RESULT(res != 0);
147
148 res = EVP_CipherFinal_ex(self->decrypt_ctx, NULL, &outlen2);
149 if (res == 0) {
150 PyErr_SetString(CryptoError, "Payload decryption failed");
151 return NULL;
152 }
153
154 return PyBytes_FromStringAndSize((const char*)self->buffer, outlen);
155}
156
157static PyObject*
158AEAD_encrypt(AEADObject *self, PyObject *args)
159{
160 const unsigned char *data, *associated;
161 Py_ssize_t data_len, associated_len;
162 int outlen, outlen2, res;
163 uint64_t pn;
164
165 if (!PyArg_ParseTuple(args, "y#y#K", &data, &data_len, &associated, &associated_len, &pn))
166 return NULL;
167
168 if (data_len > PACKET_LENGTH_MAX) {
169 PyErr_SetString(CryptoError, "Invalid payload length");
170 return NULL;
171 }
172
173 memcpy(self->nonce, self->iv, AEAD_NONCE_LENGTH);
174 for (int i = 0; i < 8; ++i) {
175 self->nonce[AEAD_NONCE_LENGTH - 1 - i] ^= (uint8_t)(pn >> 8 * i);
176 }
177
178 res = EVP_CipherInit_ex(self->encrypt_ctx, NULL, NULL, self->key, self->nonce, 1);
179 CHECK_RESULT(res != 0);
180
181 res = EVP_CipherUpdate(self->encrypt_ctx, NULL, &outlen, associated, associated_len);
182 CHECK_RESULT(res != 0);
183
184 res = EVP_CipherUpdate(self->encrypt_ctx, self->buffer, &outlen, data, data_len);
185 CHECK_RESULT(res != 0);
186
187 res = EVP_CipherFinal_ex(self->encrypt_ctx, NULL, &outlen2);
188 CHECK_RESULT(res != 0 && outlen2 == 0);
189
190 res = EVP_CIPHER_CTX_ctrl(self->encrypt_ctx, EVP_CTRL_CCM_GET_TAG, AEAD_TAG_LENGTH, self->buffer + outlen);
191 CHECK_RESULT(res != 0);
192
193 return PyBytes_FromStringAndSize((const char*)self->buffer, outlen + AEAD_TAG_LENGTH);
194}
195
196static PyMethodDef AEAD_methods[] = {
197 {"decrypt", (PyCFunction)AEAD_decrypt, METH_VARARGS, ""},
198 {"encrypt", (PyCFunction)AEAD_encrypt, METH_VARARGS, ""},
199
200 {NULL}
201};
202
203static PyType_Slot AEADType_slots[] = {
204 {Py_tp_dealloc, AEAD_dealloc},
205 {Py_tp_methods, AEAD_methods},
206 {Py_tp_doc, "AEAD objects"},
207 {Py_tp_init, AEAD_init},
208 {0, 0},
209};
210
211static PyType_Spec AEADType_spec = {
212 MODULE_NAME ".AEADType",
213 sizeof(AEADObject),
214 0,
215 Py_TPFLAGS_DEFAULT,
216 AEADType_slots
217};
218
219/* HeaderProtection */
220
221typedef struct {
222 PyObject_HEAD
223 EVP_CIPHER_CTX *ctx;
224 int is_chacha20;
225 unsigned char buffer[PACKET_LENGTH_MAX];
226 unsigned char mask[31];
227 unsigned char zero[5];
228} HeaderProtectionObject;
229
230static PyObject *HeaderProtectionType;
231
232static int
233HeaderProtection_init(HeaderProtectionObject *self, PyObject *args, PyObject *kwargs)
234{
235 const char *cipher_name;
236 const unsigned char *key;
237 Py_ssize_t cipher_name_len, key_len;
238 int res;
239
240 if (!PyArg_ParseTuple(args, "y#y#", &cipher_name, &cipher_name_len, &key, &key_len))
241 return -1;
242
243 const EVP_CIPHER *evp_cipher = EVP_get_cipherbyname(cipher_name);
244 if (evp_cipher == 0) {
245 PyErr_Format(CryptoError, "Invalid cipher name: %s", cipher_name);
246 return -1;
247 }
248
249 memset(self->mask, 0, sizeof(self->mask));
250 memset(self->zero, 0, sizeof(self->zero));
251 self->is_chacha20 = cipher_name_len == 8 && memcmp(cipher_name, "chacha20", 8) == 0;
252
253 self->ctx = EVP_CIPHER_CTX_new();
254 CHECK_RESULT_CTOR(self->ctx != 0);
255
256 res = EVP_CipherInit_ex(self->ctx, evp_cipher, NULL, NULL, NULL, 1);
257 CHECK_RESULT_CTOR(res != 0);
258
259 res = EVP_CIPHER_CTX_set_key_length(self->ctx, key_len);
260 CHECK_RESULT_CTOR(res != 0);
261
262 res = EVP_CipherInit_ex(self->ctx, NULL, NULL, key, NULL, 1);
263 CHECK_RESULT_CTOR(res != 0);
264
265 return 0;
266}
267
268static void
269HeaderProtection_dealloc(HeaderProtectionObject *self)
270{
271 EVP_CIPHER_CTX_free(self->ctx);
272 PyTypeObject *tp = Py_TYPE(self);
273 freefunc free = PyType_GetSlot(tp, Py_tp_free);
274 free(self);
275 Py_DECREF(tp);
276}
277
278static int HeaderProtection_mask(HeaderProtectionObject *self, const unsigned char* sample)
279{
280 int outlen;
281 if (self->is_chacha20) {
282 return EVP_CipherInit_ex(self->ctx, NULL, NULL, NULL, sample, 1) &&
283 EVP_CipherUpdate(self->ctx, self->mask, &outlen, self->zero, sizeof(self->zero));
284 } else {
285 return EVP_CipherUpdate(self->ctx, self->mask, &outlen, sample, SAMPLE_LENGTH);
286 }
287}
288
289static PyObject*
290HeaderProtection_apply(HeaderProtectionObject *self, PyObject *args)
291{
292 const unsigned char *header, *payload;
293 Py_ssize_t header_len, payload_len;
294 int res;
295
296 if (!PyArg_ParseTuple(args, "y#y#", &header, &header_len, &payload, &payload_len))
297 return NULL;
298
299 int pn_length = (header[0] & 0x03) + 1;
300 int pn_offset = header_len - pn_length;
301
302 res = HeaderProtection_mask(self, payload + PACKET_NUMBER_LENGTH_MAX - pn_length);
303 CHECK_RESULT(res != 0);
304
305 memcpy(self->buffer, header, header_len);
306 memcpy(self->buffer + header_len, payload, payload_len);
307
308 if (self->buffer[0] & 0x80) {
309 self->buffer[0] ^= self->mask[0] & 0x0F;
310 } else {
311 self->buffer[0] ^= self->mask[0] & 0x1F;
312 }
313
314 for (int i = 0; i < pn_length; ++i) {
315 self->buffer[pn_offset + i] ^= self->mask[1 + i];
316 }
317
318 return PyBytes_FromStringAndSize((const char*)self->buffer, header_len + payload_len);
319}
320
321static PyObject*
322HeaderProtection_remove(HeaderProtectionObject *self, PyObject *args)
323{
324 const unsigned char *packet;
325 Py_ssize_t packet_len;
326 int pn_offset, res;
327
328 if (!PyArg_ParseTuple(args, "y#I", &packet, &packet_len, &pn_offset))
329 return NULL;
330
331 res = HeaderProtection_mask(self, packet + pn_offset + PACKET_NUMBER_LENGTH_MAX);
332 CHECK_RESULT(res != 0);
333
334 memcpy(self->buffer, packet, pn_offset + PACKET_NUMBER_LENGTH_MAX);
335
336 if (self->buffer[0] & 0x80) {
337 self->buffer[0] ^= self->mask[0] & 0x0F;
338 } else {
339 self->buffer[0] ^= self->mask[0] & 0x1F;
340 }
341
342 int pn_length = (self->buffer[0] & 0x03) + 1;
343 uint32_t pn_truncated = 0;
344 for (int i = 0; i < pn_length; ++i) {
345 self->buffer[pn_offset + i] ^= self->mask[1 + i];
346 pn_truncated = self->buffer[pn_offset + i] | (pn_truncated << 8);
347 }
348
349 return Py_BuildValue("y#i", self->buffer, pn_offset + pn_length, pn_truncated);
350}
351
352static PyMethodDef HeaderProtection_methods[] = {
353 {"apply", (PyCFunction)HeaderProtection_apply, METH_VARARGS, ""},
354 {"remove", (PyCFunction)HeaderProtection_remove, METH_VARARGS, ""},
355 {NULL}
356};
357
358static PyType_Slot HeaderProtectionType_slots[] = {
359 {Py_tp_dealloc, HeaderProtection_dealloc},
360 {Py_tp_methods, HeaderProtection_methods},
361 {Py_tp_doc, "HeaderProtection objects"},
362 {Py_tp_init, HeaderProtection_init},
363 {0, 0},
364};
365
366static PyType_Spec HeaderProtectionType_spec = {
367 MODULE_NAME ".HeaderProtectionType",
368 sizeof(HeaderProtectionObject),
369 0,
370 Py_TPFLAGS_DEFAULT,
371 HeaderProtectionType_slots
372};
373
374static struct PyModuleDef moduledef = {
375 PyModuleDef_HEAD_INIT,
376 MODULE_NAME, /* m_name */
377 "Cryptography utilities.", /* m_doc */
378 -1, /* m_size */
379 NULL, /* m_methods */
380 NULL, /* m_reload */
381 NULL, /* m_traverse */
382 NULL, /* m_clear */
383 NULL, /* m_free */
384};
385
386PyMODINIT_FUNC
387PyInit__crypto(void)
388{
389 PyObject* m;
390
391 m = PyModule_Create(&moduledef);
392 if (m == NULL)
393 return NULL;
394
395 CryptoError = PyErr_NewException(MODULE_NAME ".CryptoError", PyExc_ValueError, NULL);
396 Py_INCREF(CryptoError);
397 PyModule_AddObject(m, "CryptoError", CryptoError);
398
399 AEADType = PyType_FromSpec(&AEADType_spec);
400 if (AEADType == NULL)
401 return NULL;
402 PyModule_AddObject(m, "AEAD", AEADType);
403
404 HeaderProtectionType = PyType_FromSpec(&HeaderProtectionType_spec);
405 if (HeaderProtectionType == NULL)
406 return NULL;
407 PyModule_AddObject(m, "HeaderProtection", HeaderProtectionType);
408
409 // ensure required ciphers are initialised
410 EVP_add_cipher(EVP_aes_128_ecb());
411 EVP_add_cipher(EVP_aes_128_gcm());
412 EVP_add_cipher(EVP_aes_256_ecb());
413 EVP_add_cipher(EVP_aes_256_gcm());
414
415 return m;
416}
417