comparison libervia/backend/plugins/plugin_sec_aesgcm.py @ 4270:0d7bb4df2343

Reformatted code base using black.
author Goffi <goffi@goffi.org>
date Wed, 19 Jun 2024 18:44:57 +0200
parents 4b842c1fb686
children
comparison
equal deleted inserted replaced
4269:64a85ce8be70 4270:0d7bb4df2343
43 C.PI_TYPE: "SEC", 43 C.PI_TYPE: "SEC",
44 C.PI_PROTOCOLS: ["OMEMO Media sharing"], 44 C.PI_PROTOCOLS: ["OMEMO Media sharing"],
45 C.PI_DEPENDENCIES: ["XEP-0363", "XEP-0384", "DOWNLOAD", "ATTACH"], 45 C.PI_DEPENDENCIES: ["XEP-0363", "XEP-0384", "DOWNLOAD", "ATTACH"],
46 C.PI_MAIN: "AESGCM", 46 C.PI_MAIN: "AESGCM",
47 C.PI_HANDLER: "no", 47 C.PI_HANDLER: "no",
48 C.PI_DESCRIPTION: dedent(_("""\ 48 C.PI_DESCRIPTION: dedent(
49 _(
50 """\
49 Implementation of AES-GCM scheme, a way to encrypt files (not official XMPP standard). 51 Implementation of AES-GCM scheme, a way to encrypt files (not official XMPP standard).
50 See https://xmpp.org/extensions/inbox/omemo-media-sharing.html for details 52 See https://xmpp.org/extensions/inbox/omemo-media-sharing.html for details
51 """)), 53 """
54 )
55 ),
52 } 56 }
53 57
54 AESGCM_RE = re.compile( 58 AESGCM_RE = re.compile(
55 r'aesgcm:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9' 59 r"aesgcm:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9"
56 r'()@:%_\+.~#?&\/\/=]*)') 60 r"()@:%_\+.~#?&\/\/=]*)"
61 )
57 62
58 63
59 class AESGCM(object): 64 class AESGCM(object):
60 65
61 def __init__(self, host): 66 def __init__(self, host):
62 self.host = host 67 self.host = host
63 log.info(_("AESGCM plugin initialization")) 68 log.info(_("AESGCM plugin initialization"))
64 self._http_upload = host.plugins['XEP-0363'] 69 self._http_upload = host.plugins["XEP-0363"]
65 self._attach = host.plugins["ATTACH"] 70 self._attach = host.plugins["ATTACH"]
66 host.plugins["DOWNLOAD"].register_scheme( 71 host.plugins["DOWNLOAD"].register_scheme("aesgcm", self.download)
67 "aesgcm", self.download 72 self._attach.register(self.can_handle_attachment, self.attach, encrypted=True)
68 )
69 self._attach.register(
70 self.can_handle_attachment, self.attach, encrypted=True)
71 host.trigger.add("XEP-0363_upload_pre_slot", self._upload_pre_slot) 73 host.trigger.add("XEP-0363_upload_pre_slot", self._upload_pre_slot)
72 host.trigger.add("XEP-0363_upload", self._upload_trigger) 74 host.trigger.add("XEP-0363_upload", self._upload_trigger)
73 host.trigger.add("message_received", self._message_received_trigger) 75 host.trigger.add("message_received", self._message_received_trigger)
74 76
75 async def download(self, client, uri_parsed, dest_path, options): 77 async def download(self, client, uri_parsed, dest_path, options):
82 iv_size = 16 84 iv_size = 16
83 elif len(fragment) == 44: 85 elif len(fragment) == 44:
84 iv_size = 12 86 iv_size = 12
85 else: 87 else:
86 raise ValueError( 88 raise ValueError(
87 f"Invalid URL fragment, can't decrypt file at {uri_parsed.get_url()}") 89 f"Invalid URL fragment, can't decrypt file at {uri_parsed.get_url()}"
90 )
88 91
89 iv, key = fragment[:iv_size], fragment[iv_size:] 92 iv, key = fragment[:iv_size], fragment[iv_size:]
90 93
91 decryptor = ciphers.Cipher( 94 decryptor = ciphers.Cipher(
92 ciphers.algorithms.AES(key), 95 ciphers.algorithms.AES(key),
93 modes.GCM(iv), 96 modes.GCM(iv),
94 backend=backends.default_backend(), 97 backend=backends.default_backend(),
95 ).decryptor() 98 ).decryptor()
96 99
97 download_url = parse.urlunparse( 100 download_url = parse.urlunparse(
98 ('https', uri_parsed.netloc, uri_parsed.path, '', '', '')) 101 ("https", uri_parsed.netloc, uri_parsed.path, "", "", "")
99 102 )
100 if options.get('ignore_tls_errors', False): 103
101 log.warning( 104 if options.get("ignore_tls_errors", False):
102 "TLS certificate check disabled, this is highly insecure" 105 log.warning("TLS certificate check disabled, this is highly insecure")
103 )
104 treq_client = treq_client_no_ssl 106 treq_client = treq_client_no_ssl
105 else: 107 else:
106 treq_client = treq 108 treq_client = treq
107 109
108 head_data = await treq_client.head(download_url) 110 head_data = await treq_client.head(download_url)
109 content_length = int(head_data.headers.getRawHeaders('content-length')[0]) 111 content_length = int(head_data.headers.getRawHeaders("content-length")[0])
110 # the 128 bits tag is put at the end 112 # the 128 bits tag is put at the end
111 file_size = content_length - 16 113 file_size = content_length - 16
112 114
113 file_obj = stream.SatFile( 115 file_obj = stream.SatFile(
114 self.host, 116 self.host,
115 client, 117 client,
116 dest_path, 118 dest_path,
117 mode="wb", 119 mode="wb",
118 size = file_size, 120 size=file_size,
119 ) 121 )
120 122
121 progress_id = file_obj.uid 123 progress_id = file_obj.uid
122 124
123 resp = await treq_client.get(download_url, unbuffered=True) 125 resp = await treq_client.get(download_url, unbuffered=True)
124 if resp.code == 200: 126 if resp.code == 200:
125 d = treq.collect(resp, partial( 127 d = treq.collect(
126 self.on_data_download, 128 resp,
127 client=client, 129 partial(
128 file_obj=file_obj, 130 self.on_data_download,
129 decryptor=decryptor)) 131 client=client,
132 file_obj=file_obj,
133 decryptor=decryptor,
134 ),
135 )
130 else: 136 else:
131 d = defer.Deferred() 137 d = defer.Deferred()
132 self.host.plugins["DOWNLOAD"].errback_download(file_obj, d, resp) 138 self.host.plugins["DOWNLOAD"].errback_download(file_obj, d, resp)
133 return progress_id, d 139 return progress_id, d
134 140
139 return False 145 return False
140 else: 146 else:
141 return True 147 return True
142 148
143 async def _upload_cb(self, client, filepath, filename, extra): 149 async def _upload_cb(self, client, filepath, filename, extra):
144 extra['encryption'] = C.ENC_AES_GCM 150 extra["encryption"] = C.ENC_AES_GCM
145 return await self._http_upload.file_http_upload( 151 return await self._http_upload.file_http_upload(
146 client=client, 152 client=client, filepath=filepath, filename=filename, extra=extra
147 filepath=filepath,
148 filename=filename,
149 extra=extra
150 ) 153 )
151 154
152 async def attach(self, client, data): 155 async def attach(self, client, data):
153 # XXX: the attachment removal/resend code below is due to the one file per 156 # XXX: the attachment removal/resend code below is due to the one file per
154 # message limitation of OMEMO media sharing unofficial XEP. We have to remove 157 # message limitation of OMEMO media sharing unofficial XEP. We have to remove
155 # attachments from original message, and send them one by one. 158 # attachments from original message, and send them one by one.
156 # TODO: this is to be removed when a better mechanism is available with OMEMO (now 159 # TODO: this is to be removed when a better mechanism is available with OMEMO (now
157 # possible with the 0.4 version of OMEMO, it's possible to encrypt other stanza 160 # possible with the 0.4 version of OMEMO, it's possible to encrypt other stanza
158 # elements than body). 161 # elements than body).
159 attachments = data["extra"][C.KEY_ATTACHMENTS] 162 attachments = data["extra"][C.KEY_ATTACHMENTS]
160 if not data['message'] or data['message'] == {'': ''}: 163 if not data["message"] or data["message"] == {"": ""}:
161 extra_attachments = attachments[1:] 164 extra_attachments = attachments[1:]
162 del attachments[1:] 165 del attachments[1:]
163 await self._attach.upload_files(client, data, upload_cb=self._upload_cb) 166 await self._attach.upload_files(client, data, upload_cb=self._upload_cb)
164 else: 167 else:
165 # we have a message, we must send first attachment separately 168 # we have a message, we must send first attachment separately
175 body_elt.addContent(attachment["url"]) 178 body_elt.addContent(attachment["url"])
176 179
177 for attachment in extra_attachments: 180 for attachment in extra_attachments:
178 # we send all remaining attachment in a separate message 181 # we send all remaining attachment in a separate message
179 await client.sendMessage( 182 await client.sendMessage(
180 to_jid=data['to'], 183 to_jid=data["to"],
181 message={'': ''}, 184 message={"": ""},
182 subject=data['subject'], 185 subject=data["subject"],
183 mess_type=data['type'], 186 mess_type=data["type"],
184 extra={C.KEY_ATTACHMENTS: [attachment]}, 187 extra={C.KEY_ATTACHMENTS: [attachment]},
185 ) 188 )
186 189
187 if ((not data['extra'] 190 if (
188 and (not data['message'] or data['message'] == {'': ''}) 191 not data["extra"]
189 and not data['subject'])): 192 and (not data["message"] or data["message"] == {"": ""})
193 and not data["subject"]
194 ):
190 # nothing left to send, we can cancel the message 195 # nothing left to send, we can cancel the message
191 raise exceptions.CancelError("Cancelled by AESGCM attachment handling") 196 raise exceptions.CancelError("Cancelled by AESGCM attachment handling")
192 197
193 def on_data_download(self, data, client, file_obj, decryptor): 198 def on_data_download(self, data, client, file_obj, decryptor):
194 if file_obj.tell() + len(data) > file_obj.size: 199 if file_obj.tell() + len(data) > file_obj.size:
220 else: 225 else:
221 decrypted = decryptor.update(data) 226 decrypted = decryptor.update(data)
222 file_obj.write(decrypted) 227 file_obj.write(decrypted)
223 228
224 def _upload_pre_slot(self, client, extra, file_metadata): 229 def _upload_pre_slot(self, client, extra, file_metadata):
225 if extra.get('encryption') != C.ENC_AES_GCM: 230 if extra.get("encryption") != C.ENC_AES_GCM:
226 return True 231 return True
227 # the tag is appended to the file 232 # the tag is appended to the file
228 file_metadata["size"] += 16 233 file_metadata["size"] += 16
229 return True 234 return True
230 235
237 ret = encryptor.finalize() 242 ret = encryptor.finalize()
238 tag = encryptor.tag 243 tag = encryptor.tag
239 return ret + tag 244 return ret + tag
240 except AlreadyFinalized: 245 except AlreadyFinalized:
241 # as we have already finalized, we can now send EOF 246 # as we have already finalized, we can now send EOF
242 return b'' 247 return b""
243 248
244 def _upload_trigger(self, client, extra, sat_file, file_producer, slot): 249 def _upload_trigger(self, client, extra, sat_file, file_producer, slot):
245 if extra.get('encryption') != C.ENC_AES_GCM: 250 if extra.get("encryption") != C.ENC_AES_GCM:
246 return True 251 return True
247 log.debug("encrypting file with AES-GCM") 252 log.debug("encrypting file with AES-GCM")
248 iv = secrets.token_bytes(12) 253 iv = secrets.token_bytes(12)
249 key = secrets.token_bytes(32) 254 key = secrets.token_bytes(32)
250 fragment = f'{iv.hex()}{key.hex()}' 255 fragment = f"{iv.hex()}{key.hex()}"
251 ori_url = parse.urlparse(slot.get) 256 ori_url = parse.urlparse(slot.get)
252 # we change the get URL with the one with aesgcm scheme and containing the 257 # we change the get URL with the one with aesgcm scheme and containing the
253 # encoded key + iv 258 # encoded key + iv
254 slot.get = parse.urlunparse(['aesgcm', *ori_url[1:5], fragment]) 259 slot.get = parse.urlunparse(["aesgcm", *ori_url[1:5], fragment])
255 260
256 # encrypted data size will be bigger than original file size 261 # encrypted data size will be bigger than original file size
257 # so we need to check with final data length to avoid a warning on close() 262 # so we need to check with final data length to avoid a warning on close()
258 sat_file.check_size_with_read = True 263 sat_file.check_size_with_read = True
259 264
268 backend=backends.default_backend(), 273 backend=backends.default_backend(),
269 ).encryptor() 274 ).encryptor()
270 275
271 if sat_file.data_cb is not None: 276 if sat_file.data_cb is not None:
272 raise exceptions.InternalError( 277 raise exceptions.InternalError(
273 f"data_cb was expected to be None, it is set to {sat_file.data_cb}") 278 f"data_cb was expected to be None, it is set to {sat_file.data_cb}"
279 )
274 280
275 # with data_cb we encrypt the file on the fly 281 # with data_cb we encrypt the file on the fly
276 sat_file.data_cb = partial(self._encrypt, encryptor=encryptor) 282 sat_file.data_cb = partial(self._encrypt, encryptor=encryptor)
277 return True 283 return True
278
279 284
280 def _pop_aesgcm_links(self, match, links): 285 def _pop_aesgcm_links(self, match, links):
281 link = match.group() 286 link = match.group()
282 if link not in links: 287 if link not in links:
283 links.append(link) 288 links.append(link)
284 return "" 289 return ""
285 290
286 def _check_aesgcm_attachments(self, client, data): 291 def _check_aesgcm_attachments(self, client, data):
287 if not data.get('message'): 292 if not data.get("message"):
288 return data 293 return data
289 links = [] 294 links = []
290 295
291 for lang, message in list(data['message'].items()): 296 for lang, message in list(data["message"].items()):
292 message = AESGCM_RE.sub( 297 message = AESGCM_RE.sub(partial(self._pop_aesgcm_links, links=links), message)
293 partial(self._pop_aesgcm_links, links=links),
294 message)
295 if links: 298 if links:
296 message = message.strip() 299 message = message.strip()
297 if not message: 300 if not message:
298 del data['message'][lang] 301 del data["message"][lang]
299 else: 302 else:
300 data['message'][lang] = message 303 data["message"][lang] = message
301 mess_encrypted = client.encryption.isEncrypted(data) 304 mess_encrypted = client.encryption.isEncrypted(data)
302 attachments = data['extra'].setdefault(C.KEY_ATTACHMENTS, []) 305 attachments = data["extra"].setdefault(C.KEY_ATTACHMENTS, [])
303 for link in links: 306 for link in links:
304 path = parse.urlparse(link).path 307 path = parse.urlparse(link).path
305 attachment = { 308 attachment = {
306 "url": link, 309 "url": link,
307 } 310 }
312 if mess_encrypted: 315 if mess_encrypted:
313 # we don't add the encrypted flag if the message itself is not 316 # we don't add the encrypted flag if the message itself is not
314 # encrypted, because the decryption key is part of the link, 317 # encrypted, because the decryption key is part of the link,
315 # so sending it over unencrypted channel is like having no 318 # so sending it over unencrypted channel is like having no
316 # encryption at all. 319 # encryption at all.
317 attachment['encrypted'] = True 320 attachment["encrypted"] = True
318 attachments.append(attachment) 321 attachments.append(attachment)
319 322
320 return data 323 return data
321 324
322 def _message_received_trigger(self, client, message_elt, post_treat): 325 def _message_received_trigger(self, client, message_elt, post_treat):