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