From e5748860e7b490076d08afe39cc564232418de70 Mon Sep 17 00:00:00 2001
From: alemi <me@alemi.dev>
Date: Mon, 15 Jul 2024 02:57:51 +0200
Subject: [PATCH] feat(cli): added cloak command to fix previous urls

---
 upub/cli/Cargo.toml   |  1 +
 upub/cli/src/cloak.rs | 46 +++++++++++++++++++++++++++++++++++++++++++
 upub/cli/src/lib.rs   | 14 ++++++++++++-
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 upub/cli/src/cloak.rs

diff --git a/upub/cli/Cargo.toml b/upub/cli/Cargo.toml
index ec11cdf..d9a0466 100644
--- a/upub/cli/Cargo.toml
+++ b/upub/cli/Cargo.toml
@@ -22,3 +22,4 @@ openssl = "0.10" # TODO handle pubkeys with a smaller crate
 clap = { version = "4.5", features = ["derive"] }
 sea-orm = { version = "0.12", features = ["macros", "sqlx-sqlite", "runtime-tokio-rustls"] }
 futures = "0.3"
+mdhtml = { path = "../../utils/mdhtml/" }
diff --git a/upub/cli/src/cloak.rs b/upub/cli/src/cloak.rs
new file mode 100644
index 0000000..7670e1e
--- /dev/null
+++ b/upub/cli/src/cloak.rs
@@ -0,0 +1,46 @@
+use futures::TryStreamExt;
+use sea_orm::{ActiveModelTrait, ActiveValue::{Set, Unchanged}, ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter, QuerySelect, SelectColumns, TransactionTrait};
+use upub::traits::{fetch::RequestError, Cloaker};
+
+pub async fn cloak(ctx: upub::Context, post_contents: bool) -> Result<(), RequestError> {
+	let tx = ctx.db().begin().await?;
+
+	{
+		let mut stream = upub::model::attachment::Entity::find()
+			.filter(upub::model::attachment::Column::Url.not_like(format!("{}%", ctx.base())))
+			.stream(ctx.db())
+			.await?;
+
+		while let Some(attachment) = stream.try_next().await? {
+			let (sig, url) = ctx.cloak(&attachment.url);
+			let mut model = attachment.into_active_model();
+			model.url = Set(upub::url!(ctx, "/proxy/{sig}/{url}"));
+			model.update(&tx).await?;
+		}
+	}
+
+	if post_contents {
+		let mut stream = upub::model::object::Entity::find()
+			.filter(upub::model::object::Column::Content.is_not_null())
+			.select_only()
+			.select_column(upub::model::object::Column::Internal)
+			.select_column(upub::model::object::Column::Content)
+			.into_tuple::<(i64, String)>()
+			.stream(ctx.db())
+			.await?;
+
+		while let Some((internal, content)) = stream.try_next().await? {
+			let sanitized = mdhtml::safe_html(&content);
+			if sanitized != content {
+				let model = upub::model::object::ActiveModel {
+					internal: Unchanged(internal),
+					content: Set(Some(sanitized)),
+					..Default::default()
+				};
+				model.update(&tx).await?;
+			}
+		}
+	}
+
+	Ok(())
+}
diff --git a/upub/cli/src/lib.rs b/upub/cli/src/lib.rs
index d07d0ba..17761b3 100644
--- a/upub/cli/src/lib.rs
+++ b/upub/cli/src/lib.rs
@@ -22,6 +22,9 @@ pub use nuke::*;
 mod thread;
 pub use thread::*;
 
+mod cloak;
+pub use cloak::*;
+
 #[derive(Debug, Clone, clap::Subcommand)]
 pub enum CliCommand {
 	/// generate fake user, note and activity
@@ -109,7 +112,14 @@ pub enum CliCommand {
 	/// attempt to fix broken threads and completely gather their context
 	Thread {
 		
-	}
+	},
+
+	/// replaces all attachment urls with proxied local versions (only useful for old instances)
+	Cloak {
+		/// also replace urls inside post contents
+		#[arg(long, default_value_t = false)]
+		post_contents: bool,
+	},
 }
 
 pub async fn run(ctx: upub::Context, command: CliCommand) -> Result<(), Box<dyn std::error::Error>> {
@@ -131,5 +141,7 @@ pub async fn run(ctx: upub::Context, command: CliCommand) -> Result<(), Box<dyn
 			Ok(nuke(ctx, for_real, delete_objects).await?),
 		CliCommand::Thread { } =>
 			Ok(thread(ctx).await?),
+		CliCommand::Cloak { post_contents } =>
+			Ok(cloak(ctx, post_contents).await?),
 	}
 }