1use std::collections::HashSet;
36use std::fmt;
37
38use age::secrecy::{ExposeSecret as _, SecretString};
39use age_core::format::{FileKey, Stanza, FILE_KEY_BYTES};
40use base64::engine::general_purpose::STANDARD_NO_PAD;
41use base64::Engine as _;
42use bech32::{FromBase32 as _, ToBase32 as _, Variant};
43use hpke::{Deserializable as _, Kem as _, Serializable as _};
44use x_wing::TryDecapsulate as _;
45use zeroize::{Zeroize as _, Zeroizing};
46
47pub use crate::members::PQ_STANZA as STANZA_TYPE;
48
49const INFO: &[u8] = b"age-encryption.org/mlkem768x25519";
51const RECIPIENT_HRP: &str = "age1pq";
53const IDENTITY_HRP: &str = "age-secret-key-pq-";
55const ENC_BYTES: usize = 1120;
57pub const LABEL: &str = "postquantum";
59
60type Kem = hpke::kem::XWing;
61type Aead = hpke::aead::ChaCha20Poly1305;
62type Kdf = hpke::kdf::HkdfSha256;
63
64struct OsRng;
67
68impl hpke::rand_core::TryRng for OsRng {
69 type Error = core::convert::Infallible;
70 fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
71 Ok(u32::from_le_bytes(crate::random_bytes::<4>()))
72 }
73 fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
74 Ok(u64::from_le_bytes(crate::random_bytes::<8>()))
75 }
76 fn try_fill_bytes(&mut self, dst: &mut [u8]) -> Result<(), Self::Error> {
77 getrandom::getrandom(dst).expect("the OS random source is available");
78 Ok(())
79 }
80}
81impl hpke::rand_core::TryCryptoRng for OsRng {}
82
83#[derive(Clone, PartialEq, Eq)]
85pub struct Recipient(<Kem as hpke::Kem>::PublicKey);
86
87impl std::str::FromStr for Recipient {
88 type Err = &'static str;
89
90 fn from_str(s: &str) -> Result<Self, Self::Err> {
91 let bytes = crate::members::pq_public_key_bytes(s).ok_or("not an age1pq1… key")?;
92 <Kem as hpke::Kem>::PublicKey::from_bytes(&bytes)
93 .map(Recipient)
94 .map_err(|_| "not a valid ML-KEM-768 + X25519 public key")
95 }
96}
97
98impl fmt::Display for Recipient {
99 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100 let s = bech32::encode(
101 RECIPIENT_HRP,
102 self.0.to_bytes().to_base32(),
103 Variant::Bech32,
104 )
105 .map_err(|_| fmt::Error)?;
106 f.write_str(&s)
107 }
108}
109
110impl fmt::Debug for Recipient {
111 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112 write!(f, "{self}")
113 }
114}
115
116impl age::Recipient for Recipient {
117 fn wrap_file_key(
118 &self,
119 file_key: &FileKey,
120 ) -> Result<(Vec<Stanza>, HashSet<String>), age::EncryptError> {
121 let (enc, body) = hpke::single_shot_seal_with_rng::<Aead, Kdf, Kem>(
122 &hpke::OpModeS::Base,
123 &self.0,
124 INFO,
125 file_key.expose_secret(),
126 &[],
127 &mut OsRng,
128 )
129 .map_err(|e| age::EncryptError::Io(std::io::Error::other(e.to_string())))?;
130 Ok((
131 vec![Stanza {
132 tag: STANZA_TYPE.to_string(),
133 args: vec![STANDARD_NO_PAD.encode(enc.to_bytes())],
134 body,
135 }],
136 HashSet::from([LABEL.to_string()]),
137 ))
138 }
139}
140
141#[derive(Clone)]
143pub struct Identity(<Kem as hpke::Kem>::PrivateKey);
144
145impl Identity {
146 pub fn generate() -> Identity {
148 let seed = Zeroizing::new(crate::random_bytes::<32>());
149 Identity::from_seed(&seed)
150 }
151
152 fn from_seed(seed: &[u8; 32]) -> Identity {
153 Identity(<Kem as hpke::Kem>::PrivateKey::from_bytes(seed).expect("32 bytes is a seed"))
154 }
155
156 pub fn to_secret_string(&self) -> SecretString {
159 let seed = Zeroizing::new(self.0.to_bytes());
160 let mut lower = bech32::encode(IDENTITY_HRP, seed.as_slice().to_base32(), Variant::Bech32)
161 .expect("the HRP is valid");
162 let upper = lower.to_uppercase();
163 lower.zeroize();
164 SecretString::from(upper)
165 }
166
167 pub fn to_public(&self) -> Recipient {
169 Recipient(Kem::sk_to_pk(&self.0))
170 }
171}
172
173impl std::str::FromStr for Identity {
174 type Err = &'static str;
175
176 fn from_str(s: &str) -> Result<Self, Self::Err> {
177 let s = s.trim();
178 if !s.starts_with("AGE-SECRET-KEY-PQ-1") {
179 return Err("not an AGE-SECRET-KEY-PQ-1… identity");
180 }
181 let (hrp, data, variant) = bech32::decode(s).map_err(|_| "invalid bech32")?;
182 if hrp != IDENTITY_HRP || variant != Variant::Bech32 {
183 return Err("not an AGE-SECRET-KEY-PQ-1… identity");
184 }
185 let bytes = Zeroizing::new(Vec::<u8>::from_base32(&data).map_err(|_| "invalid bech32")?);
186 let seed: &[u8; 32] = bytes
187 .as_slice()
188 .try_into()
189 .map_err(|_| "a post-quantum identity is 32 bytes")?;
190 Ok(Identity::from_seed(seed))
191 }
192}
193
194impl age::Identity for Identity {
195 fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, age::DecryptError>> {
196 if stanza.tag != STANZA_TYPE {
197 return None;
198 }
199 let header = || Some(Err(age::DecryptError::InvalidHeader));
200 let [arg] = &stanza.args[..] else {
203 return header();
204 };
205 if arg.len() != 1494 {
208 return header();
209 }
210 let Some(enc) = STANDARD_NO_PAD
211 .decode(arg)
212 .ok()
213 .filter(|b| b.len() == ENC_BYTES && STANDARD_NO_PAD.encode(b) == *arg)
214 else {
215 return header();
216 };
217 if stanza.body.len() != FILE_KEY_BYTES + 16 {
218 return header();
219 }
220 let seed = Zeroizing::new(<[u8; 32]>::from(self.0.to_bytes()));
223 let strict = x_wing::DecapsulationKeyRejectNonContrib::from(*seed);
224 let Ok(ct) = x_wing::Ciphertext::try_from(enc.as_slice()) else {
225 return header();
226 };
227 if strict.try_decapsulate(&ct).is_err() {
228 return header();
229 }
230 let Ok(enc) = <Kem as hpke::Kem>::EncappedKey::from_bytes(&enc) else {
231 return header();
232 };
233 match hpke::single_shot_open::<Aead, Kdf, Kem>(
234 &hpke::OpModeR::Base,
235 &self.0,
236 &enc,
237 INFO,
238 &stanza.body,
239 &[],
240 ) {
241 Err(hpke::HpkeError::DecapError) => header(),
244 Err(_) => None,
245 Ok(mut pt) => {
246 let key = FileKey::init_with_mut(|k| k.copy_from_slice(&pt));
247 pt.zeroize();
248 Some(Ok(key))
249 }
250 }
251 }
252}
253
254pub fn is_identity_file(text: &str) -> bool {
257 identity_line(text).is_some()
258}
259
260fn identity_line(text: &str) -> Option<&str> {
261 text.lines()
262 .map(str::trim)
263 .find(|l| l.starts_with("AGE-SECRET-KEY-PQ-1"))
264}
265
266pub fn parse_identity_file(text: &str) -> Result<Identity, &'static str> {
268 identity_line(text)
269 .ok_or("no AGE-SECRET-KEY-PQ-1… line")?
270 .parse()
271}
272
273pub fn identity_file(identity: &Identity, created: &str) -> Zeroizing<String> {
276 Zeroizing::new(format!(
277 "# created: {created}\n# public key: {}\n{}\n",
278 identity.to_public(),
279 identity.to_secret_string().expose_secret()
280 ))
281}