knowledge-search

commit a8d04b50d69ada55d545ca402a00b812cb590f40
parent 38026299528dd7c9fa15cef029675a08b58a8c06
Author: MTRNord <mtrnord1@gmail.com>
Date:   Fri, 24 Mar 2023 01:39:14 +0100

feat(matrix-indexer): Add writing back the access_token and device_id to the config after first login

Diffstat:
M
crates/matrix-indexer/src/config.rs
 | 
56
+++++++++++++++++++++++++++++++++++++++++++++++++++++++-
M
crates/matrix-indexer/src/main.rs
 | 
16
++++++++++++----
M
crates/matrix-indexer/src/matrix.rs
 | 
7
+++++++

3 files changed, 74 insertions(+), 5 deletions(-)
diff --git a/crates/matrix-indexer/src/config.rs b/crates/matrix-indexer/src/config.rs
@@ -1,4 +1,6 @@
-use kdl::KdlDocument;
+use std::io::Write;
+
+use kdl::{KdlDocument, KdlEntry, KdlIdentifier, KdlNode};
 use miette::{Diagnostic, NamedSource, SourceSpan};
 use thiserror::Error;
 
@@ -195,3 +197,55 @@ pub fn load() -> Config {
         }
     }
 }
+
+pub fn write_access_token(access_token: String, device_id: String) -> color_eyre::Result<()> {
+    let source = std::fs::read_to_string("config.kdl")
+        .expect("Unable to open config.kdl file. Is it present?");
+    let config = parse(&source);
+    match config {
+        Ok(mut config) => {
+            let matrix = config.get_mut("matrix");
+            if let Some(matrix) = matrix {
+                if let Some(matrix_children) = matrix.children_mut() {
+                    let matrix_nodes = matrix_children.nodes_mut();
+                    let access_token_identifier = KdlIdentifier::from("access_token");
+                    let device_id_identifier = KdlIdentifier::from("device_id");
+
+                    let mut access_token_node = KdlNode::new(access_token_identifier);
+                    let access_token_entry = KdlEntry::new(access_token);
+                    access_token_node.entries_mut().push(access_token_entry);
+
+                    let mut device_id_node = KdlNode::new(device_id_identifier);
+                    let device_id_entry = KdlEntry::new(device_id);
+                    device_id_node.entries_mut().push(device_id_entry);
+
+                    matrix_nodes.push(access_token_node);
+                    matrix_nodes.push(device_id_node);
+                    matrix_nodes.retain(|node| {
+                        node.name().value() != "username" && node.name().value() != "password"
+                    });
+
+                    matrix_nodes.sort_by(sort_by_name);
+                }
+                matrix.fmt();
+            }
+
+            let mut file = std::fs::OpenOptions::new()
+                .write(true)
+                .truncate(true)
+                .open("config.kdl")?;
+
+            file.write_all(config.to_string().as_bytes())?;
+            file.flush()?;
+            Ok(())
+        }
+        Err(error) => {
+            let err = format!("{error:?}");
+            Err(color_eyre::eyre::eyre!("{}", err))
+        }
+    }
+}
+
+fn sort_by_name(x: &KdlNode, y: &KdlNode) -> std::cmp::Ordering {
+    x.name().value().cmp(y.name().value())
+}
diff --git a/crates/matrix-indexer/src/main.rs b/crates/matrix-indexer/src/main.rs
@@ -24,9 +24,9 @@
 // I am lazy. Dont blame me!
 #![allow(missing_docs)]
 
-use color_eyre::Result;
+use color_eyre::{eyre::bail, Result};
 
-use config::load;
+use config::{load, write_access_token};
 use matrix::IndexerBot;
 
 mod config;
@@ -41,15 +41,23 @@ async fn main() -> Result<()> {
     let config = load();
 
     // TODO: config which rewrites itself to have the data after login
+    #[allow(clippy::unwrap_used)]
     let mut bot = match config.auth_data {
         config::AuthData::UsernamePassword(mxid, password) => {
-            IndexerBot::new(
+            let bot = IndexerBot::new(
                 config.homeserver_url,
                 mxid,
                 password,
                 config.indradb_endpoint,
             )
-            .await?
+            .await?;
+            let access_token = bot.access_token();
+            let device_id = bot.device_id();
+            if access_token.is_none() || device_id.is_none() {
+                bail!("Login to matrix must have failed. We got no access_token or device_id!");
+            }
+            write_access_token(access_token.unwrap(), device_id.unwrap())?;
+            bot
         }
         config::AuthData::AccessToken(mxid, access_token, device_id) => {
             IndexerBot::relogin(
diff --git a/crates/matrix-indexer/src/matrix.rs b/crates/matrix-indexer/src/matrix.rs
@@ -35,6 +35,13 @@ pub struct IndexerBot {
 }
 
 impl IndexerBot {
+    pub fn access_token(&self) -> Option<String> {
+        self.client.access_token()
+    }
+
+    pub fn device_id(&self) -> Option<String> {
+        self.client.device_id().map(ToString::to_string)
+    }
     async fn get_client(homeserver_url: String) -> Result<Client> {
         let mut client_builder = Client::builder().homeserver_url(homeserver_url);
         client_builder = client_builder.sled_store(Path::new("./matrix_data"), None)?;