upub_cli/
cloak.rs

1use futures::TryStreamExt;
2use sea_orm::{ActiveModelTrait, ActiveValue::{NotSet, Set, Unchanged}, ColumnTrait, Condition, EntityTrait, IntoActiveModel, QueryFilter, QuerySelect, SelectColumns};
3use upub::traits::{fetch::RequestError, Cloaker};
4
5pub async fn cloak(ctx: upub::Context, post_contents: bool, objects: bool, actors: bool, re_cloak: bool) -> Result<(), RequestError> {
6	let local_base = format!("{}%", ctx.base());
7	{
8		let mut select = upub::model::attachment::Entity::find();
9		if !re_cloak {
10			select = select.filter(upub::model::attachment::Column::Url.not_like(&local_base));
11		}
12
13		let mut stream = select
14			.stream(ctx.db())
15			.await?;
16
17		while let Some(attachment) = stream.try_next().await? {
18			tracing::info!("cloaking {}", attachment.url);
19			let url = ctx.cloaked(&attachment.url);
20			let mut model = attachment.into_active_model();
21			model.url = Set(url);
22			model.update(ctx.db()).await?;
23		}
24	}
25
26	if objects {
27		let mut select = upub::model::object::Entity::find()
28			.filter(upub::model::object::Column::Image.is_not_null());
29
30		if !re_cloak {
31			select = select.filter(upub::model::object::Column::Image.not_like(&local_base));
32		}
33		
34		let mut stream = select
35			.select_only()
36			.select_column(upub::model::object::Column::Internal)
37			.select_column(upub::model::object::Column::Image)
38			.into_tuple::<(i64, String)>()
39			.stream(ctx.db())
40			.await?;
41
42		while let Some((internal, image)) = stream.try_next().await? {
43			tracing::info!("cloaking object image {image}");
44			let model = upub::model::object::ActiveModel {
45				internal: Unchanged(internal),
46				image: Set(Some(ctx.cloaked(&image))),
47				..Default::default()
48			};
49			model.update(ctx.db()).await?;
50		}
51	}
52
53	if actors {
54		let mut select = upub::model::actor::Entity::find();
55
56		if !re_cloak {
57			select = select
58				.filter(
59					Condition::any()
60						.add(upub::model::actor::Column::Image.not_like(&local_base))
61						.add(upub::model::actor::Column::Icon.not_like(&local_base))
62				);
63		}
64
65		let mut stream = select
66			.select_only()
67			.select_column(upub::model::actor::Column::Internal)
68			.select_column(upub::model::actor::Column::Image)
69			.select_column(upub::model::actor::Column::Icon)
70			.into_tuple::<(i64, Option<String>, Option<String>)>()
71			.stream(ctx.db())
72			.await?;
73
74		while let Some((internal, image, icon)) = stream.try_next().await? {
75			tracing::info!("cloaking user #{internal}");
76			if image.is_none() && icon.is_none() { continue }
77			// TODO can this if/else/else be made nicer??
78			let image = if let Some(img) = image {
79				if !img.starts_with(ctx.base()) {
80					Set(Some(ctx.cloaked(&img)))
81				} else {
82					NotSet
83				}
84			} else {
85				NotSet
86			};
87			let icon = if let Some(icn) = icon {
88				if !icn.starts_with(ctx.base()) {
89					Set(Some(ctx.cloaked(&icn)))
90				} else {
91					NotSet
92				}
93			} else {
94				NotSet
95			};
96			let model = upub::model::actor::ActiveModel {
97				internal: Unchanged(internal),
98				image, icon,
99				..Default::default()
100			};
101			model.update(ctx.db()).await?;
102		}
103	}
104
105	if post_contents {
106		let mut stream = upub::model::object::Entity::find()
107			.filter(upub::model::object::Column::Content.like("%<img%"))
108			.select_only()
109			.select_column(upub::model::object::Column::Internal)
110			.select_column(upub::model::object::Column::Content)
111			.into_tuple::<(i64, String)>()
112			.stream(ctx.db())
113			.await?;
114
115		while let Some((internal, content)) = stream.try_next().await? {
116			let sanitized = ctx.sanitize(&content);
117			if sanitized != content {
118				tracing::info!("sanitizing object #{internal}");
119				let model = upub::model::object::ActiveModel {
120					internal: Unchanged(internal),
121					content: Set(Some(sanitized)),
122					..Default::default()
123				};
124				model.update(ctx.db()).await?;
125			}
126		}
127	}
128
129	Ok(())
130}