1use std::{
2 collections::{HashMap, HashSet},
3 io,
4};
5
6use age::{Identity, Recipient};
7use age_core::format::{FileKey, Stanza};
8use age_plugin::{
9 identity::{self, IdentityPluginV1},
10 recipient::{self, RecipientPluginV1},
11 Callbacks, PluginHandler,
12};
13use bincode::{config, Decode, Encode};
14
15use tlock_age::{internal::STANZA_TAG, Header};
16
17pub const ROUND_ENV: &str = "ROUND";
19
20#[derive(Debug, Encode, Decode, PartialEq, Clone)]
21pub struct RecipientInfo {
27 hash: Vec<u8>,
28 public_key_bytes: Vec<u8>,
29 genesis_time: u64,
30 period: u64,
31}
32
33impl RecipientInfo {
34 pub fn new(hash: &[u8], public_key_bytes: &[u8], genesis_time: u64, period: u64) -> Self {
35 Self {
36 hash: hash.to_vec(),
37 public_key_bytes: public_key_bytes.to_vec(),
38 genesis_time,
39 period,
40 }
41 }
42
43 fn serialize(&self) -> Vec<u8> {
44 bincode::encode_to_vec(self, config::standard()).unwrap()
45 }
46
47 fn deserialize(data: &[u8]) -> Self {
48 let (result, _) = bincode::decode_from_slice(data, config::standard()).unwrap();
49 result
50 }
51
52 pub fn hash(&self) -> Vec<u8> {
53 self.hash.clone()
54 }
55 pub fn public_key_bytes(&self) -> Vec<u8> {
56 self.public_key_bytes.clone()
57 }
58 pub fn genesis_time(&self) -> u64 {
59 self.genesis_time
60 }
61 pub fn period(&self) -> u64 {
62 self.period
63 }
64}
65
66pub struct RecipientPlugin {
67 plugin_name: String,
68 info: Option<RecipientInfo>,
69 parse_round: fn(&RecipientInfo, &str) -> u64,
70}
71
72impl RecipientPlugin {
73 pub fn new(plugin_name: &str, parse_round: fn(&RecipientInfo, &str) -> u64) -> Self {
74 Self {
75 plugin_name: plugin_name.to_owned(),
76 info: None,
77 parse_round,
78 }
79 }
80
81 pub fn plugin_name(&self) -> String {
82 self.plugin_name.clone()
83 }
84
85 pub fn info(&self) -> Option<RecipientInfo> {
86 self.info.clone()
87 }
88
89 pub fn parse_round(&self, round: &str) -> u64 {
90 (self.parse_round)(&self.info().unwrap(), round)
91 }
92}
93
94impl RecipientPluginV1 for RecipientPlugin {
95 fn add_recipient(
96 &mut self,
97 index: usize,
98 plugin_name: &str,
99 bytes: &[u8],
100 ) -> Result<(), recipient::Error> {
101 if plugin_name == self.plugin_name() {
102 let chain = RecipientInfo::deserialize(bytes);
103 self.info = Some(chain);
104 Ok(())
105 } else {
106 Err(recipient::Error::Recipient {
107 index,
108 message: "unsupported plugin".to_owned(),
109 })
110 }
111 }
112
113 fn add_identity(
114 &mut self,
115 _index: usize,
116 _plugin_name: &str,
117 _bytes: &[u8],
118 ) -> Result<(), recipient::Error> {
119 todo!()
120 }
121
122 fn labels(&mut self) -> HashSet<String> {
123 HashSet::new()
124 }
125
126 fn wrap_file_keys(
127 &mut self,
128 file_keys: Vec<FileKey>,
129 mut callbacks: impl Callbacks<recipient::Error>,
130 ) -> io::Result<Result<Vec<Vec<Stanza>>, Vec<recipient::Error>>> {
131 let round = if let Ok(round) = std::env::var(ROUND_ENV) {
132 round
133 } else {
134 let prompt_message = "Enter decryption round: ";
135 match callbacks.request_public(prompt_message) {
136 Ok(round) => round.unwrap_or("".to_owned()),
137 Err(err) => return Err(err),
138 }
139 };
140 let round = self.parse_round(&round);
141
142 let info = self.info().unwrap();
143
144 let recipient =
145 tlock_age::internal::Recipient::new(&info.hash, &info.public_key_bytes, round);
146 Ok(Ok(file_keys
147 .into_iter()
148 .map(|file_key| {
149 let (stanzas, _labels) = recipient.wrap_file_key(&file_key).unwrap();
150 stanzas
151 })
152 .collect()))
153 }
154}
155
156pub enum IdentityFormat {
160 RAW,
161 HTTP,
162}
163
164#[derive(Debug, Encode, Decode, PartialEq, Clone)]
165pub enum IdentityInfo {
167 RawIdentityInfo(RawIdentityInfo),
168 HTTPIdentityInfo(HTTPIdentityInfo),
169}
170
171impl IdentityInfo {
172 fn serialize(&self) -> Vec<u8> {
173 bincode::encode_to_vec(self, config::standard()).unwrap()
174 }
175
176 fn deserialize(data: &[u8]) -> Self {
177 let (result, _) = bincode::decode_from_slice(data, config::standard()).unwrap();
178 result
179 }
180
181 pub fn format(&self) -> IdentityFormat {
182 match self {
183 Self::RawIdentityInfo(_) => IdentityFormat::RAW,
184 Self::HTTPIdentityInfo(_) => IdentityFormat::HTTP,
185 }
186 }
187}
188
189impl From<RawIdentityInfo> for IdentityInfo {
190 fn from(value: RawIdentityInfo) -> Self {
191 IdentityInfo::RawIdentityInfo(value)
192 }
193}
194
195impl From<HTTPIdentityInfo> for IdentityInfo {
196 fn from(value: HTTPIdentityInfo) -> Self {
197 IdentityInfo::HTTPIdentityInfo(value)
198 }
199}
200
201#[derive(Debug, Encode, Decode, PartialEq, Clone)]
202pub struct RawIdentityInfo {
203 signature: Vec<u8>,
204}
205
206impl RawIdentityInfo {
207 pub fn new(signature: &[u8]) -> Self {
208 Self {
209 signature: signature.to_vec(),
210 }
211 }
212}
213
214#[derive(Debug, Encode, Decode, PartialEq, Clone)]
215pub struct HTTPIdentityInfo {
216 url: String,
217}
218
219impl HTTPIdentityInfo {
220 pub fn new(url: &str) -> Self {
221 Self {
222 url: url.to_owned(),
223 }
224 }
225}
226
227pub struct IdentityPlugin {
228 plugin_name: String,
229 info: Option<IdentityInfo>,
230 get_signature: fn(url: &str, header: &Header) -> Vec<u8>,
231}
232
233impl IdentityPlugin {
234 pub fn new(
235 plugin_name: &str,
236 get_signature: fn(url: &str, header: &Header) -> Vec<u8>,
237 ) -> Self {
238 Self {
239 plugin_name: plugin_name.to_owned(),
240 info: None,
241 get_signature,
242 }
243 }
244}
245
246impl IdentityPluginV1 for IdentityPlugin {
247 fn add_identity(
248 &mut self,
249 index: usize,
250 plugin_name: &str,
251 bytes: &[u8],
252 ) -> Result<(), identity::Error> {
253 if plugin_name == self.plugin_name.as_str() {
254 let info = IdentityInfo::deserialize(bytes);
255 self.info = Some(info);
256 Ok(())
257 } else {
258 Err(identity::Error::Identity {
259 index,
260 message: "unsupported plugin".to_owned(),
261 })
262 }
263 }
264
265 fn unwrap_file_keys(
266 &mut self,
267 files: Vec<Vec<Stanza>>,
268 _callbacks: impl Callbacks<identity::Error>,
269 ) -> io::Result<HashMap<usize, Result<FileKey, Vec<identity::Error>>>> {
270 let mut file_keys = HashMap::with_capacity(files.len());
271
272 for (file, stanzas) in files.iter().enumerate() {
273 for stanza in stanzas.iter() {
274 if stanza.tag != STANZA_TAG {
275 continue;
276 }
277 if stanza.args.len() != 2 {
278 continue; }
280 let [round, hash] = [stanza.args[0].clone(), stanza.args[1].clone()];
281 let round = round.parse().unwrap();
282 let hash = hex::decode(hash).unwrap();
283 let header = Header::new(round, &hash);
284
285 let signature = match self.info.as_ref().unwrap() {
286 IdentityInfo::HTTPIdentityInfo(info) => {
287 (self.get_signature)(info.url.as_str(), &header)
288 }
289 IdentityInfo::RawIdentityInfo(info) => info.signature.clone(),
290 };
291 let identity = tlock_age::internal::Identity::new(&hash, &signature);
292
293 let file_key = identity.unwrap_stanza(stanza).unwrap();
294 let r = file_key.map_err(|e| {
295 vec![identity::Error::Identity {
296 index: file,
297 message: format!("{e}"),
298 }]
299 });
300
301 file_keys.entry(file).or_insert_with(|| r);
302 }
303 }
304 Ok(file_keys)
305 }
306}
307
308pub struct TlockPluginHandler {
310 plugin_name: String,
311 parse_round: fn(&RecipientInfo, &str) -> u64,
312 get_signature: fn(&str, &Header) -> Vec<u8>,
313}
314
315impl TlockPluginHandler {
316 pub fn new(
317 plugin_name: &str,
318 parse_round: fn(&RecipientInfo, &str) -> u64,
319 get_signature: fn(&str, &Header) -> Vec<u8>,
320 ) -> Self {
321 Self {
322 plugin_name: plugin_name.to_owned(),
323 parse_round,
324 get_signature,
325 }
326 }
327}
328
329impl PluginHandler for TlockPluginHandler {
330 type RecipientV1 = RecipientPlugin;
331 type IdentityV1 = IdentityPlugin;
332
333 fn recipient_v1(self) -> io::Result<Self::RecipientV1> {
334 Ok(RecipientPlugin::new(&self.plugin_name, self.parse_round))
335 }
336
337 fn identity_v1(self) -> io::Result<Self::IdentityV1> {
338 Ok(IdentityPlugin::new(&self.plugin_name, self.get_signature))
339 }
340}
341
342pub fn run_state_machine(
345 state_machine: String,
346 plugin_name: &str,
347 parse_round: fn(&RecipientInfo, &str) -> u64,
348 get_signature: fn(&str, &Header) -> Vec<u8>,
349) -> io::Result<()> {
350 age_plugin::run_state_machine(
352 &state_machine,
353 TlockPluginHandler::new(plugin_name, parse_round, get_signature),
354 )
355}
356
357pub fn print_new_identity(plugin_name: &str, identity: &IdentityInfo, recipient: &RecipientInfo) {
359 age_plugin::print_new_identity(plugin_name, &identity.serialize(), &recipient.serialize())
360}