Skip to content

Commit f113746

Browse files
nnathannnathan
authored
Implement Armored Ciphertext (#116)
1 parent e7e51fe commit f113746

File tree

8 files changed

+118
-51
lines changed

8 files changed

+118
-51
lines changed

‎Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ crate-type = ["cdylib"]
1818

1919
[dependencies]
2020
age-core = "0.11"
21-
age = { version = "0.11.1", features = ["ssh", "plugin"] }
21+
age = { version = "0.11.1", features = ["ssh", "plugin", "armor"] }
2222
pyo3 = { version = "0.24.2", features = [
2323
"extension-module",
2424
"abi3",

‎dev-requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,4 @@ build
22
wheel
33
twine
44
maturin
5+
parameterized

‎pyrage-stubs/pyrage/__init__.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,9 @@ class DecryptError(Exception):
4242
...
4343

4444

45-
def encrypt(plaintext: bytes, recipients: Sequence[_Recipient]) -> bytes: ...
46-
def encrypt_file(infile: str, outfile: str, recipients: Sequence[_Recipient]) -> None: ...
47-
def encrypt_io(in_io: BufferedIOBase, out_io: BufferedIOBase, recipients: Sequence[_Recipient]) -> bytes: ...
45+
def encrypt(plaintext: bytes, recipients: Sequence[_Recipient], armored: bool) -> bytes: ...
46+
def encrypt_file(infile: str, outfile: str, recipients: Sequence[_Recipient], armored: bool) -> None: ...
47+
def encrypt_io(in_io: BufferedIOBase, out_io: BufferedIOBase, recipients: Sequence[_Recipient], armored: bool) -> bytes: ...
4848

4949
def decrypt(ciphertext: bytes, identities: Sequence[_Identity]) -> bytes: ...
5050
def decrypt_file(infile: str, outfile: str, identities: Sequence[_Identity]) -> None: ...

‎pyrage-stubs/pyrage/passphrase.pyi‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,2 @@
1-
def encrypt(plaintext: bytes, passphrase: str) -> bytes:
2-
...
3-
4-
5-
def decrypt(ciphertext: bytes, passphrase: str) -> bytes:
6-
...
1+
def encrypt(plaintext: bytes, passphrase: str, armored: bool = False) -> bytes: ...
2+
def decrypt(ciphertext: bytes, passphrase: str, armored: bool = False) -> bytes: ...

‎src/lib.rs‎

Lines changed: 55 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ use std::io::Write;
55
use std::{fs::File, io::Read};
66

77
use age::{
8-
DecryptError as RageDecryptError, EncryptError as RageEncryptError, Encryptor, Identity,
9-
Recipient,
8+
armor::ArmoredReader, armor::ArmoredWriter, armor::Format, DecryptError as RageDecryptError,
9+
EncryptError as RageEncryptError, Encryptor, Identity, Recipient,
1010
};
1111
use age_core::format::{FileKey, Stanza};
1212
use pyo3::{
@@ -136,10 +136,12 @@ impl<'source> FromPyObject<'source> for Box<dyn PyrageIdentity> {
136136
create_exception!(pyrage, EncryptError, PyException);
137137

138138
#[pyfunction]
139+
#[pyo3(signature = (plaintext, recipients, armored=false))]
139140
fn encrypt<'p>(
140141
py: Python<'p>,
141142
plaintext: &[u8],
142143
recipients: Vec<Box<dyn PyrageRecipient>>,
144+
armored: bool,
143145
) -> PyResult<Bound<'p, PyBytes>> {
144146
// This turns each `dyn PyrageRecipient` into a `dyn Recipient`, which
145147
// is what the underlying `age` API expects.
@@ -151,13 +153,25 @@ fn encrypt<'p>(
151153
let encryptor = Encryptor::with_recipients(recipients.iter().map(|r| r.as_ref()))
152154
.map_err(|_| EncryptError::new_err("expected at least one recipient"))?;
153155
let mut encrypted = vec![];
154-
let mut writer = encryptor
155-
.wrap_output(&mut encrypted)
156-
.map_err(|e| EncryptError::new_err(e.to_string()))?;
156+
157+
let mut writer = match armored {
158+
true => encryptor
159+
.wrap_output(ArmoredWriter::wrap_output(
160+
&mut encrypted,
161+
Format::AsciiArmor,
162+
)?)
163+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
164+
false => encryptor
165+
.wrap_output(ArmoredWriter::wrap_output(&mut encrypted, Format::Binary)?)
166+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
167+
};
168+
157169
writer
158170
.write_all(plaintext)
159171
.map_err(|e| EncryptError::new_err(e.to_string()))?;
160172
writer
173+
.finish()
174+
.map_err(|e| EncryptError::new_err(e.to_string()))?
161175
.finish()
162176
.map_err(|e| EncryptError::new_err(e.to_string()))?;
163177

@@ -166,10 +180,12 @@ fn encrypt<'p>(
166180
}
167181

168182
#[pyfunction]
183+
#[pyo3(signature = (infile, outfile, recipients, armored=false))]
169184
fn encrypt_file(
170185
infile: String,
171186
outfile: String,
172187
recipients: Vec<Box<dyn PyrageRecipient>>,
188+
armored: bool,
173189
) -> PyResult<()> {
174190
// This turns each `dyn PyrageRecipient` into a `dyn Recipient`, which
175191
// is what the underlying `age` API expects.
@@ -186,13 +202,21 @@ fn encrypt_file(
186202

187203
let encryptor = Encryptor::with_recipients(recipients.iter().map(|r| r.as_ref()))
188204
.map_err(|_| EncryptError::new_err("expected at least one recipient"))?;
189-
let mut writer = encryptor
190-
.wrap_output(&mut writer)
191-
.map_err(|e| EncryptError::new_err(e.to_string()))?;
205+
206+
let mut writer = match armored {
207+
true => encryptor
208+
.wrap_output(ArmoredWriter::wrap_output(&mut writer, Format::AsciiArmor)?)
209+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
210+
false => encryptor
211+
.wrap_output(ArmoredWriter::wrap_output(&mut writer, Format::Binary)?)
212+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
213+
};
192214

193215
std::io::copy(&mut reader, &mut writer).map_err(|e| EncryptError::new_err(e.to_string()))?;
194216

195217
writer
218+
.finish()
219+
.map_err(|e| EncryptError::new_err(e.to_string()))?
196220
.finish()
197221
.map_err(|e| EncryptError::new_err(e.to_string()))?;
198222

@@ -209,8 +233,8 @@ fn decrypt<'p>(
209233
) -> PyResult<Bound<'p, PyBytes>> {
210234
let identities = identities.iter().map(|pi| pi.as_ref().as_identity());
211235

212-
let decryptor =
213-
age::Decryptor::new(ciphertext).map_err(|e| DecryptError::new_err(e.to_string()))?;
236+
let decryptor = age::Decryptor::new(ArmoredReader::new(ciphertext))
237+
.map_err(|e| DecryptError::new_err(e.to_string()))?;
214238

215239
let mut decrypted = vec![];
216240
let mut reader = decryptor
@@ -238,8 +262,8 @@ fn decrypt_file(
238262
let reader = std::io::BufReader::new(reader);
239263
let mut writer = std::io::BufWriter::new(writer);
240264

241-
let decryptor =
242-
age::Decryptor::new_buffered(reader).map_err(|e| DecryptError::new_err(e.to_string()))?;
265+
let decryptor = age::Decryptor::new_buffered(ArmoredReader::new(reader))
266+
.map_err(|e| DecryptError::new_err(e.to_string()))?;
243267

244268
let mut reader = decryptor
245269
.decrypt(identities)
@@ -256,10 +280,12 @@ fn from_pyobject(file: PyObject, read_only: bool) -> PyResult<PyFileLikeObject>
256280
}
257281

258282
#[pyfunction]
283+
#[pyo3(signature = (reader, writer, recipients, armored=false))]
259284
fn encrypt_io(
260285
reader: PyObject,
261286
writer: PyObject,
262287
recipients: Vec<Box<dyn PyrageRecipient>>,
288+
armored: bool,
263289
) -> PyResult<()> {
264290
// This turns each `dyn PyrageRecipient` into a `dyn Recipient`, which
265291
// is what the underlying `age` API expects.
@@ -271,15 +297,27 @@ fn encrypt_io(
271297
let writer = from_pyobject(writer, false)?;
272298
let mut reader = std::io::BufReader::new(reader);
273299
let mut writer = std::io::BufWriter::new(writer);
300+
274301
let encryptor = Encryptor::with_recipients(recipients.iter().map(|r| r.as_ref()))
275302
.map_err(|_| EncryptError::new_err("expected at least one recipient"))?;
276-
let mut writer = encryptor
277-
.wrap_output(&mut writer)
278-
.map_err(|e| EncryptError::new_err(e.to_string()))?;
303+
304+
let mut writer = match armored {
305+
true => encryptor
306+
.wrap_output(ArmoredWriter::wrap_output(&mut writer, Format::AsciiArmor)?)
307+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
308+
false => encryptor
309+
.wrap_output(ArmoredWriter::wrap_output(&mut writer, Format::Binary)?)
310+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
311+
};
312+
279313
std::io::copy(&mut reader, &mut writer).map_err(|e| EncryptError::new_err(e.to_string()))?;
314+
280315
writer
316+
.finish()
317+
.map_err(|e| EncryptError::new_err(e.to_string()))?
281318
.finish()
282319
.map_err(|e| EncryptError::new_err(e.to_string()))?;
320+
283321
Ok(())
284322
}
285323

@@ -294,8 +332,8 @@ fn decrypt_io(
294332
let writer = from_pyobject(writer, false)?;
295333
let reader = std::io::BufReader::new(reader);
296334
let mut writer = std::io::BufWriter::new(writer);
297-
let decryptor =
298-
age::Decryptor::new_buffered(reader).map_err(|e| DecryptError::new_err(e.to_string()))?;
335+
let decryptor = age::Decryptor::new_buffered(ArmoredReader::new(reader))
336+
.map_err(|e| DecryptError::new_err(e.to_string()))?;
299337
let mut reader = decryptor
300338
.decrypt(identities)
301339
.map_err(|e| DecryptError::new_err(e.to_string()))?;

‎src/passphrase.rs‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,22 +3,44 @@ use std::{
33
iter,
44
};
55

6-
use age::{scrypt, Decryptor, Encryptor};
6+
use age::{
7+
armor::ArmoredReader, armor::ArmoredWriter, armor::Format, scrypt, Decryptor, Encryptor,
8+
};
79
use pyo3::{prelude::*, types::PyBytes};
810

911
use crate::{DecryptError, EncryptError};
1012

1113
#[pyfunction]
12-
fn encrypt<'p>(py: Python<'p>, plaintext: &[u8], passphrase: &str) -> PyResult<Bound<'p, PyBytes>> {
14+
#[pyo3(signature = (plaintext, passphrase, armored=false))]
15+
fn encrypt<'p>(
16+
py: Python<'p>,
17+
plaintext: &[u8],
18+
passphrase: &str,
19+
armored: bool,
20+
) -> PyResult<Bound<'p, PyBytes>> {
1321
let encryptor = Encryptor::with_user_passphrase(passphrase.into());
1422
let mut encrypted = vec![];
15-
let mut writer = encryptor
16-
.wrap_output(&mut encrypted)
17-
.map_err(|e| EncryptError::new_err(e.to_string()))?;
23+
24+
let writer_result = match armored {
25+
true => encryptor.wrap_output(
26+
ArmoredWriter::wrap_output(&mut encrypted, Format::AsciiArmor)
27+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
28+
),
29+
false => encryptor.wrap_output(
30+
ArmoredWriter::wrap_output(&mut encrypted, Format::Binary)
31+
.map_err(|e| EncryptError::new_err(e.to_string()))?,
32+
),
33+
};
34+
35+
let mut writer = writer_result.map_err(|e| EncryptError::new_err(e.to_string()))?;
36+
1837
writer
1938
.write_all(plaintext)
2039
.map_err(|e| EncryptError::new_err(e.to_string()))?;
40+
2141
writer
42+
.finish()
43+
.map_err(|e| EncryptError::new_err(e.to_string()))?
2244
.finish()
2345
.map_err(|e| EncryptError::new_err(e.to_string()))?;
2446

@@ -31,7 +53,8 @@ fn decrypt<'p>(
3153
ciphertext: &[u8],
3254
passphrase: &str,
3355
) -> PyResult<Bound<'p, PyBytes>> {
34-
let decryptor = Decryptor::new(ciphertext).map_err(|e| DecryptError::new_err(e.to_string()))?;
56+
let decryptor = Decryptor::new_buffered(ArmoredReader::new(ciphertext))
57+
.map_err(|e| DecryptError::new_err(e.to_string()))?;
3558
let mut decrypted = vec![];
3659
let mut reader = decryptor
3760
.decrypt(iter::once(&scrypt::Identity::new(passphrase.into()) as _))

‎test/test_passphrase.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
import unittest
22

3+
from parameterized import parameterized
4+
35
from pyrage import passphrase
46

57

68
class TestPassphrase(unittest.TestCase):
7-
def test_roundtrip(self):
9+
@parameterized.expand([(False,), (True,)])
10+
def test_roundtrip(self, armored):
811
plaintext = b"junk"
9-
encrypted = passphrase.encrypt(plaintext, "some password")
12+
encrypted = passphrase.encrypt(plaintext, "some password", armored=armored)
1013
decrypted = passphrase.decrypt(encrypted, "some password")
1114

1215
self.assertEqual(plaintext, decrypted)

0 commit comments

Comments
 (0)