From ef7fc3c6a90bdb1937588d7e2b60259b59d1d9a0 Mon Sep 17 00:00:00 2001
From: networkException <git@nwex.de>
Date: Fri, 18 Oct 2024 19:38:41 +0200
Subject: [PATCH] Main: Add support for fetching all media from a element chat
 export

---
 src/main.rs | 62 +++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 53 insertions(+), 9 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index b74c5d7..4e63da3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,6 @@
-use std::{collections::HashSet, fs};
+use std::path::PathBuf;
+use std::collections::HashSet;
+use std::fs;
 use std::process::exit;
 use std::io::Read;
 
@@ -9,7 +11,8 @@ use log::{info, debug, error};
 use serde::Deserialize;
 use rustyline_async::{Readline, ReadlineEvent};
 
-use matrix_sdk::{matrix_auth::{MatrixSession, MatrixSessionTokens}, Client, SessionMeta};
+use matrix_sdk::{Client, SessionMeta};
+use matrix_sdk::matrix_auth::{MatrixSession, MatrixSessionTokens};
 use matrix_sdk::ruma::events::room::{MediaSource, message::{MessageType, RoomMessageEventContent}};
 use matrix_sdk::ruma::api::client::error::{ErrorBody, ErrorKind};
 use matrix_sdk::ruma::api::client::{authenticated_media, media};
@@ -139,7 +142,7 @@ async fn supports_authenticated_media(client: &Client) -> Result<bool> {
     })
 }
 
-async fn process_media(source: MediaSource, body: Option<String>, server: Option<String>) -> Result<()> {
+async fn process_media(source: MediaSource, body: Option<String>, server: &Option<String>) -> Result<()> {
     let uri = match source.clone() {
         MediaSource::Plain(uri) => uri,
         MediaSource::Encrypted(encrypted) => encrypted.url,
@@ -239,6 +242,10 @@ struct Arguments {
     #[clap(long, env)]
     server: Option<String>,
 
+    /// The path to a JSON export from Element's export feature
+    #[clap(long, env)]
+    element_export_path: Option<PathBuf>,
+
     mxc_url: Option<String>,
 }
 
@@ -249,10 +256,47 @@ async fn main() -> Result<()> {
     let Arguments {
         server,
         mxc_url,
+        element_export_path
     } = Arguments::parse();
 
-    if mxc_url.is_some() {
-        if let Err(error) = process_media(MediaSource::Plain(OwnedMxcUri::from(mxc_url.unwrap())), None, server).await {
+    if let Some(element_export_path) = element_export_path {
+        let file_contents = fs::read_to_string(element_export_path).expect("to be able to read the element export json");
+        let file_json: serde_json::Value = serde_json::from_str(&file_contents).expect("to be able to parse the json");
+
+        let events = file_json
+            .as_object().expect("the top level json to be an object")
+            .get("messages").expect("the top level json object to have a property 'messages'")
+            .as_array().expect("the messages property to be an array");
+
+        for event in events {
+            let event_object = event.as_object().expect("event to be an object");
+            let event_type = event_object.get("type").expect("event to have a type")
+                .as_str().expect("event type to be a string");
+
+            if event_type != "m.room.message" {
+                continue;
+            }
+
+            let message_event: DecryptedEventSource = serde_json::from_value(event.clone()).expect("event to be parsable");
+
+            debug!("Parsed event with {} content", message_event.content.msgtype.msgtype());
+
+            if let Err(error) = match message_event.content.msgtype {
+                MessageType::Audio(content) => process_media(content.source, Some(content.body), &server).await,
+                MessageType::File(content) => process_media(content.source, Some(content.body), &server).await,
+                MessageType::Image(content) => process_media(content.source, Some(content.body), &server).await,
+                MessageType::Video(content) => process_media(content.source, Some(content.body), &server).await,
+                _ => continue
+            } {
+                error!("Unable to process media: {}", error);
+            }
+        }
+
+        return Ok(());
+    }
+
+    if let Some(mxc_url) = mxc_url {
+        if let Err(error) = process_media(MediaSource::Plain(OwnedMxcUri::from(mxc_url)), None, &server).await {
             error!("Unable to process media: {}", error);
         }
 
@@ -305,10 +349,10 @@ async fn main() -> Result<()> {
     debug!("Parsed event with {} content", event.content.msgtype.msgtype());
 
     if let Err(error) = match event.content.msgtype {
-        MessageType::Audio(content) => process_media(content.source, Some(content.body), server).await,
-        MessageType::File(content) => process_media(content.source, Some(content.body), server).await,
-        MessageType::Image(content) => process_media(content.source, Some(content.body), server).await,
-        MessageType::Video(content) => process_media(content.source, Some(content.body), server).await,
+        MessageType::Audio(content) => process_media(content.source, Some(content.body), &server).await,
+        MessageType::File(content) => process_media(content.source, Some(content.body), &server).await,
+        MessageType::Image(content) => process_media(content.source, Some(content.body), &server).await,
+        MessageType::Video(content) => process_media(content.source, Some(content.body), &server).await,
         _ => {
             error!("Event content is not known to contain media, exiting");
             exit(1);