Skip to main content

age_plugin_tlock/
lib.rs

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
17/// Environment variable read to get round information non-interactively.
18pub const ROUND_ENV: &str = "ROUND";
19
20#[derive(Debug, Encode, Decode, PartialEq, Clone)]
21/// Recipient information as defined for the age-plugin-tlock
22/// These are required to encrypt information offline
23/// hash is required for the stanza
24/// public_key_bytes for encrypting towards
25/// genesis_time and period to parse round information
26pub 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
156/// Identity format as defined for the age-plugin-tlock
157/// RAW allows for offline decryption of a specific round
158/// HTTP allows for online decryption of an arbitrary round
159pub enum IdentityFormat {
160    RAW,
161    HTTP,
162}
163
164#[derive(Debug, Encode, Decode, PartialEq, Clone)]
165/// Identity information as defined for the age-plugin-tlock
166pub 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; // TODO: should be an error
279                }
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
308/// Plugin handler for age-plugin-tlock
309pub 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
342/// Run the state machine for the plugin, as defined on [GitHub](https://github.com/C2SP/C2SP/blob/main/age-plugin.md).
343/// This is the entry point for the plugin. It is called by the age client.
344pub 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    // The plugin was started by an age client; run the state machine.
351    age_plugin::run_state_machine(
352        &state_machine,
353        TlockPluginHandler::new(plugin_name, parse_round, get_signature),
354    )
355}
356
357/// Print the new identity information.
358pub fn print_new_identity(plugin_name: &str, identity: &IdentityInfo, recipient: &RecipientInfo) {
359    age_plugin::print_new_identity(plugin_name, &identity.serialize(), &recipient.serialize())
360}