From 13f8a8214dffa043f641772160ef404846bb2b23 Mon Sep 17 00:00:00 2001
From: r0r-5chach <r0r-5chach.xyz@proton.me>
Date: Tue, 11 Jun 2024 18:26:50 +0100
Subject: [PATCH] added restore handler

---
 Cargo.lock     | 31 +++++++++++++++++++------------
 Cargo.toml     |  2 +-
 src/account.rs | 41 +++++++++++++++++++++++++++++++++--------
 3 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 75a4816..4bceabb 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -364,13 +364,13 @@ dependencies = [
  "anyhow",
  "axum",
  "chrono",
+ "futures",
  "http",
  "mongodb",
  "pyo3",
  "serde",
  "serde_json",
  "tokio",
- "tokio-stream",
  "vrd 0.0.7",
 ]
 
@@ -435,6 +435,21 @@ version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c"
 
+[[package]]
+name = "futures"
+version = "0.3.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-io",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+]
+
 [[package]]
 name = "futures-channel"
 version = "0.3.30"
@@ -442,6 +457,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78"
 dependencies = [
  "futures-core",
+ "futures-sink",
 ]
 
 [[package]]
@@ -496,9 +512,11 @@ version = "0.3.30"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48"
 dependencies = [
+ "futures-channel",
  "futures-core",
  "futures-io",
  "futures-macro",
+ "futures-sink",
  "futures-task",
  "memchr",
  "pin-project-lite",
@@ -1703,17 +1721,6 @@ dependencies = [
  "tokio",
 ]
 
-[[package]]
-name = "tokio-stream"
-version = "0.1.15"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "267ac89e0bec6e691e5813911606935d77c476ff49024f98abcea3e7b15e37af"
-dependencies = [
- "futures-core",
- "pin-project-lite",
- "tokio",
-]
-
 [[package]]
 name = "tokio-util"
 version = "0.7.11"
diff --git a/Cargo.toml b/Cargo.toml
index 3e73cd7..65af6e3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,11 +9,11 @@ edition = "2021"
 anyhow = "1.0.86"
 axum = "0.7.5"
 chrono = "0.4.38"
+futures = "0.3.30"
 http = "1.1.0"
 mongodb = { version = "2.8.2", features = ["bson-chrono-0_4", "tokio-runtime"]}
 pyo3 = { version = "0.21.2", features = ["auto-initialize"]}
 serde = "1.0.203"
 serde_json = "1.0.117"
 tokio = "1.38.0"
-tokio-stream = "0.1.15"
 vrd = "0.0.7"
diff --git a/src/account.rs b/src/account.rs
index 6932178..dc65d39 100644
--- a/src/account.rs
+++ b/src/account.rs
@@ -1,9 +1,12 @@
 pub mod db;
 
+use std::collections::VecDeque;
+
 use anyhow::Result;
 use axum::{
     body::Body, response::{IntoResponse, Response}, routing::{get, post}, Json, Router
 };
+use futures::{stream::{self, StreamExt, TryStreamExt}, FutureExt};
 use serde::{Deserialize, Serialize};
 use crate::ApiResult;
 use db::{get_database, get_users, LogEntry, Mole, User};
@@ -160,12 +163,36 @@ async fn post_sign_out(headers: HeaderMap) -> ApiResult {
         },
     }
 }
-async fn get_restore() {
-    //TODO: Restore
-    //find all moles that match user_id
-    //find all logs that match mole_ids
-    //Return as Json 
+async fn get_restore(headers: HeaderMap) -> ApiResult {
+   let api = headers["api_key"].to_str()?;
+   let db = &get_database().await?;
 
+   match auth(api).await? {
+       Some(user) => {
+           let moles: Vec<Mole> = db.collection::<Mole>("moles")
+               .find(doc! { "_user_id": user.id }, None).await?.try_collect().await?;
+           
+           let logs: Vec<LogEntry> = stream::iter(&moles)
+               .then(|mole| async move {
+                   db.collection::<LogEntry>("logs")
+                       .find(doc! { "_mole_id": mole.id}, None)
+                       .await.unwrap()
+                       .try_collect()
+                       .await.unwrap_or_else(|_| vec![])
+               })
+               .flat_map(stream::iter).collect().await;
+
+           Ok(Response::builder()
+               .status(StatusCode::OK)
+               .body(Json(json!({
+                   "_moles": moles,
+                   "_logs": logs
+               })).into_response().into_body()).unwrap())
+       },
+       None => Ok(Response::builder()
+                .status(StatusCode::UNAUTHORIZED)
+                .body(Body::from("API Key is incorrect"))?),
+   }
 }
 
 async fn post_backup(headers: HeaderMap, Json(user_data): Json<BackupData>) -> ApiResult {
@@ -189,10 +216,8 @@ async fn post_backup(headers: HeaderMap, Json(user_data): Json<BackupData>) -> A
 async fn backup_docs_if_new<T: serde::Serialize>(docs: Vec<T>, collection: &'static str) -> Result<()> {
     let _ = docs.iter()
         .map(|doc| to_document(doc).unwrap())
-        .map(|doc| {
-            tokio::spawn(async move {
+        .map(|doc| async move {
                 let _ = insert_if_new(doc, collection).await;
-            })
         });
 
     Ok(())