Selaa lähdekoodia

add encrypted audio handler

subDesTagesMitExtraKaese 2 vuotta sitten
vanhempi
commit
57ad32d0b5
1 muutettua tiedostoa jossa 21 lisäystä ja 4 poistoa
  1. 21 4
      main.py

+ 21 - 4
main.py

@@ -19,22 +19,39 @@ creds = botlib.Creds(
 config = botlib.Config()
 config.encryption_enabled = True
 config.emoji_verify = False
-config.ignore_unverified_devices = False
+config.ignore_unverified_devices = True
 config.store_path = '/data/crypto_store/'
 bot = botlib.Bot(creds, config)
 
 asr = ASR(os.getenv('ASR_MODEL', 'tiny'))
 
 @bot.listener.on_custom_event(nio.RoomMessageAudio)
-async def on_audio_message(room, event):
+async def on_message_audio(room, event):
+  await on_audio(room, event, False)
+
+@bot.listener.on_custom_event(nio.RoomEncryptedAudio)
+async def on_encrypted_audio(room, event):
+  await on_audio(room, event, True)
+
+async def on_audio(room, event, encrypted):
   print(room.machine_name, event.sender, event.body, event.url)
   match = botlib.MessageMatch(room, event, bot)
   if match.is_not_from_this_bot():
     await bot.async_client.room_typing(room.machine_name, True, timeout=120000)
     url = urlparse(event.url)
     response = await bot.async_client.download(server_name=url.netloc, media_id=url.path[1:])
+    if encrypted:
+      print("decrypting...")
+      data = nio.crypto.attachments.decrypt_attachment(
+        response.body,
+        event.source["content"]["file"]["key"]["k"],
+        event.source["content"]["file"]["hashes"]["sha256"],
+        event.source["content"]["file"]["iv"],
+      )
+    else:
+      data = response.body
     print(response)
-    result = await asr.transcribe(response.body)
+    result = await asr.transcribe(data)
 
     await bot.async_client.room_typing(room.machine_name, False)
     if response.filename:
@@ -50,4 +67,4 @@ async def on_audio_message(room, event):
 
 if __name__ == "__main__":
   asr.load_model()
-  bot.run()
+  bot.run()