Mercurial > libervia-backend
diff 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 |
line wrap: on
line diff
--- a/libervia/backend/plugins/plugin_sec_aesgcm.py Tue Jun 18 12:06:45 2024 +0200 +++ b/libervia/backend/plugins/plugin_sec_aesgcm.py Wed Jun 19 18:44:57 2024 +0200 @@ -45,15 +45,20 @@ C.PI_DEPENDENCIES: ["XEP-0363", "XEP-0384", "DOWNLOAD", "ATTACH"], C.PI_MAIN: "AESGCM", C.PI_HANDLER: "no", - C.PI_DESCRIPTION: dedent(_("""\ + C.PI_DESCRIPTION: dedent( + _( + """\ Implementation of AES-GCM scheme, a way to encrypt files (not official XMPP standard). See https://xmpp.org/extensions/inbox/omemo-media-sharing.html for details - """)), + """ + ) + ), } AESGCM_RE = re.compile( - r'aesgcm:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9' - r'()@:%_\+.~#?&\/\/=]*)') + r"aesgcm:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9" + r"()@:%_\+.~#?&\/\/=]*)" +) class AESGCM(object): @@ -61,13 +66,10 @@ def __init__(self, host): self.host = host log.info(_("AESGCM plugin initialization")) - self._http_upload = host.plugins['XEP-0363'] + self._http_upload = host.plugins["XEP-0363"] self._attach = host.plugins["ATTACH"] - host.plugins["DOWNLOAD"].register_scheme( - "aesgcm", self.download - ) - self._attach.register( - self.can_handle_attachment, self.attach, encrypted=True) + host.plugins["DOWNLOAD"].register_scheme("aesgcm", self.download) + self._attach.register(self.can_handle_attachment, self.attach, encrypted=True) host.trigger.add("XEP-0363_upload_pre_slot", self._upload_pre_slot) host.trigger.add("XEP-0363_upload", self._upload_trigger) host.trigger.add("message_received", self._message_received_trigger) @@ -84,7 +86,8 @@ iv_size = 12 else: raise ValueError( - f"Invalid URL fragment, can't decrypt file at {uri_parsed.get_url()}") + f"Invalid URL fragment, can't decrypt file at {uri_parsed.get_url()}" + ) iv, key = fragment[:iv_size], fragment[iv_size:] @@ -95,18 +98,17 @@ ).decryptor() download_url = parse.urlunparse( - ('https', uri_parsed.netloc, uri_parsed.path, '', '', '')) + ("https", uri_parsed.netloc, uri_parsed.path, "", "", "") + ) - if options.get('ignore_tls_errors', False): - log.warning( - "TLS certificate check disabled, this is highly insecure" - ) + if options.get("ignore_tls_errors", False): + log.warning("TLS certificate check disabled, this is highly insecure") treq_client = treq_client_no_ssl else: treq_client = treq head_data = await treq_client.head(download_url) - content_length = int(head_data.headers.getRawHeaders('content-length')[0]) + content_length = int(head_data.headers.getRawHeaders("content-length")[0]) # the 128 bits tag is put at the end file_size = content_length - 16 @@ -115,18 +117,22 @@ client, dest_path, mode="wb", - size = file_size, + size=file_size, ) progress_id = file_obj.uid resp = await treq_client.get(download_url, unbuffered=True) if resp.code == 200: - d = treq.collect(resp, partial( - self.on_data_download, - client=client, - file_obj=file_obj, - decryptor=decryptor)) + d = treq.collect( + resp, + partial( + self.on_data_download, + client=client, + file_obj=file_obj, + decryptor=decryptor, + ), + ) else: d = defer.Deferred() self.host.plugins["DOWNLOAD"].errback_download(file_obj, d, resp) @@ -141,12 +147,9 @@ return True async def _upload_cb(self, client, filepath, filename, extra): - extra['encryption'] = C.ENC_AES_GCM + extra["encryption"] = C.ENC_AES_GCM return await self._http_upload.file_http_upload( - client=client, - filepath=filepath, - filename=filename, - extra=extra + client=client, filepath=filepath, filename=filename, extra=extra ) async def attach(self, client, data): @@ -157,7 +160,7 @@ # possible with the 0.4 version of OMEMO, it's possible to encrypt other stanza # elements than body). attachments = data["extra"][C.KEY_ATTACHMENTS] - if not data['message'] or data['message'] == {'': ''}: + if not data["message"] or data["message"] == {"": ""}: extra_attachments = attachments[1:] del attachments[1:] await self._attach.upload_files(client, data, upload_cb=self._upload_cb) @@ -177,16 +180,18 @@ for attachment in extra_attachments: # we send all remaining attachment in a separate message await client.sendMessage( - to_jid=data['to'], - message={'': ''}, - subject=data['subject'], - mess_type=data['type'], + to_jid=data["to"], + message={"": ""}, + subject=data["subject"], + mess_type=data["type"], extra={C.KEY_ATTACHMENTS: [attachment]}, ) - if ((not data['extra'] - and (not data['message'] or data['message'] == {'': ''}) - and not data['subject'])): + if ( + not data["extra"] + and (not data["message"] or data["message"] == {"": ""}) + and not data["subject"] + ): # nothing left to send, we can cancel the message raise exceptions.CancelError("Cancelled by AESGCM attachment handling") @@ -222,7 +227,7 @@ file_obj.write(decrypted) def _upload_pre_slot(self, client, extra, file_metadata): - if extra.get('encryption') != C.ENC_AES_GCM: + if extra.get("encryption") != C.ENC_AES_GCM: return True # the tag is appended to the file file_metadata["size"] += 16 @@ -239,19 +244,19 @@ return ret + tag except AlreadyFinalized: # as we have already finalized, we can now send EOF - return b'' + return b"" def _upload_trigger(self, client, extra, sat_file, file_producer, slot): - if extra.get('encryption') != C.ENC_AES_GCM: + if extra.get("encryption") != C.ENC_AES_GCM: return True log.debug("encrypting file with AES-GCM") iv = secrets.token_bytes(12) key = secrets.token_bytes(32) - fragment = f'{iv.hex()}{key.hex()}' + fragment = f"{iv.hex()}{key.hex()}" ori_url = parse.urlparse(slot.get) # we change the get URL with the one with aesgcm scheme and containing the # encoded key + iv - slot.get = parse.urlunparse(['aesgcm', *ori_url[1:5], fragment]) + slot.get = parse.urlunparse(["aesgcm", *ori_url[1:5], fragment]) # encrypted data size will be bigger than original file size # so we need to check with final data length to avoid a warning on close() @@ -270,13 +275,13 @@ if sat_file.data_cb is not None: raise exceptions.InternalError( - f"data_cb was expected to be None, it is set to {sat_file.data_cb}") + f"data_cb was expected to be None, it is set to {sat_file.data_cb}" + ) # with data_cb we encrypt the file on the fly sat_file.data_cb = partial(self._encrypt, encryptor=encryptor) return True - def _pop_aesgcm_links(self, match, links): link = match.group() if link not in links: @@ -284,22 +289,20 @@ return "" def _check_aesgcm_attachments(self, client, data): - if not data.get('message'): + if not data.get("message"): return data links = [] - for lang, message in list(data['message'].items()): - message = AESGCM_RE.sub( - partial(self._pop_aesgcm_links, links=links), - message) + for lang, message in list(data["message"].items()): + message = AESGCM_RE.sub(partial(self._pop_aesgcm_links, links=links), message) if links: message = message.strip() if not message: - del data['message'][lang] + del data["message"][lang] else: - data['message'][lang] = message + data["message"][lang] = message mess_encrypted = client.encryption.isEncrypted(data) - attachments = data['extra'].setdefault(C.KEY_ATTACHMENTS, []) + attachments = data["extra"].setdefault(C.KEY_ATTACHMENTS, []) for link in links: path = parse.urlparse(link).path attachment = { @@ -314,7 +317,7 @@ # encrypted, because the decryption key is part of the link, # so sending it over unencrypted channel is like having no # encryption at all. - attachment['encrypted'] = True + attachment["encrypted"] = True attachments.append(attachment) return data