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 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}