Repository: woodruffw/pyrage (language: python)

Issue:
Encryption and decryption
We need APIs for encryption and decryption.

Rough sketch:

```python
def encrypt(plaintext: bytes, recipients: Sequence[Recipient]) -> bytes:
    ...

def decrypt(ciphertext: bytes, identity: Identity) -> bytes:
   ...
```

That doesn't cover ASCII armoring, which could be an optional `kwarg` on each.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..ffd1688
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,41 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Format
+        run: cargo fmt && git diff --exit-code
+      - name: Lint
+        run: |
+          rustup update
+          rustup component add clippy
+          cargo clippy -- -D warnings
+  test:
+    strategy:
+      matrix:
+        python:
+          - "3.7"
+          - "3.8"
+          - "3.9"
+          - "3.10"
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python }}
+
+    - name: Build
+      run: make develop
+
+    - name: Test
+      run: make test
diff --git a/Cargo.toml b/Cargo.toml
index 12fdf3c..748bb89 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,5 +25,6 @@ requires-python = ">=3.7"
 crate-type = ["cdylib"]
 
 [dependencies]
+age-core = "0.8"
 age = { version = "0.8", features = ["ssh"] }
 pyo3 = { version = "0.16", features = ["extension-module"] }
diff --git a/Makefile b/Makefile
index 6907bb4..1947163 100644
--- a/Makefile
+++ b/Makefile
@@ -12,8 +12,12 @@ env/pyvenv.cfg:
 
 .PHONY: develop
 develop: env/pyvenv.cfg
-	./env/bin/maturin develop
+	. ./env/bin/activate && maturin develop
 
 .PHONY: build
 build: env/pyvenv.cfg
-	./env/bin/maturin build
+	. ./env/bin/activate && maturin build
+
+.PHONY: test
+test: develop
+	./env/bin/python -m unittest
diff --git a/src/lib.rs b/src/lib.rs
index 2fddc0a..1ce4256 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,175 @@
-use pyo3::prelude::*;
+use std::io::{Read, Write};
+
+use age::{DecryptError, EncryptError, Encryptor, Identity, Recipient};
+use age_core::format::{FileKey, Stanza};
+use pyo3::{
+    exceptions::{PyTypeError, PyValueError},
+    prelude::*,
+    types::PyBytes,
+};
 
 mod x25519;
 
+// This is a wrapper trait for age's `Recipient`, providing trait downcasting.
+//
+// We need this so that we can pass multiple different types of recipients
+// into the Python-level `encrypt` API.
+trait PyrageRecipient: Recipient {
+    fn as_recipient(self: Box<Self>) -> Box<dyn Recipient>;
+}
+
+// This is a wrapper trait for age's `Identity`, providing trait downcasting.
+//
+// We need this so that we can pass multiple different types of identities
+// into the Python-level `decrypt` API.
+trait PyrageIdentity: Identity {
+    fn as_identity(&self) -> &dyn Identity;
+}
+
+// This macro generates two trait impls for each passed in type:
+//
+// * An age `Receipient` impl, using the underlying trait impl.
+// * A `PyrageRecipient` impl, by consuming the instance and downcasting.
+macro_rules! recipient_traits {
+    ($($t:ty),+) => {
+        $(
+            impl Recipient for $t {
+                fn wrap_file_key(&self, file_key: &FileKey) -> Result<Vec<Stanza>, EncryptError> {
+                    self.0.wrap_file_key(file_key)
+                }
+            }
+
+            impl PyrageRecipient for $t {
+                fn as_recipient(self: Box<Self>) -> Box<dyn Recipient> {
+                    self as Box<dyn Recipient>
+                }
+            }
+        )*
+    }
+}
+
+recipient_traits!(x25519::Recipient);
+
+// This macro generates two trait impls for each passed in type:
+//
+// * An age `Identity` impl, using the underlying trait impl.
+// * A `PyrageIdentity` impl, by borrowing the instance and downcasting.
+macro_rules! identity_traits {
+    ($($t:ty),+) => {
+        $(
+            impl Identity for $t {
+                fn unwrap_stanza(&self, stanza: &Stanza) -> Option<Result<FileKey, DecryptError>> {
+                    self.0.unwrap_stanza(stanza)
+                }
+            }
+
+            impl PyrageIdentity for $t {
+                fn as_identity(&self) -> &dyn Identity {
+                    self as &dyn Identity
+                }
+            }
+        )*
+    }
+}
+
+identity_traits!(x25519::Identity);
+
+// This is where the magic happens, and why we need to do the trait dance
+// above: `FromPyObject` is a third-party trait, so we need to implement it
+// for `Box<dyn PyrageRecipient>` instead of `Box<dyn Recipient>`.
+//
+// The implementation itself is straightforward: we try to turn the
+// `PyAny` into each concrete recipient type, which we then perform the trait
+// cast on.
+impl<'source> FromPyObject<'source> for Box<dyn PyrageRecipient> {
+    fn extract(ob: &'source PyAny) -> PyResult<Self> {
+        if let Ok(recipient) = ob.extract::<x25519::Recipient>() {
+            Ok(Box::new(recipient) as Box<dyn PyrageRecipient>)
+        } else {
+            Err(PyTypeError::new_err(
+                "invalid type (expected a recipient type)",
+            ))
+        }
+    }
+}
+
+// Similar to the above: we try to turn the `PyAny` into a concrete identity type,
+// which we then perform the trait cast on.
+impl<'source> FromPyObject<'source> for Box<dyn PyrageIdentity> {
+    fn extract(ob: &'source PyAny) -> PyResult<Self> {
+        if let Ok(identity) = ob.extract::<x25519::Identity>() {
+            Ok(Box::new(identity) as Box<dyn PyrageIdentity>)
+        } else {
+            Err(PyTypeError::new_err(
+                "invalid type (expected an identity type)",
+            ))
+        }
+    }
+}
+
+#[pyfunction]
+fn encrypt<'p>(
+    py: Python<'p>,
+    plaintext: &[u8],
+    recipients: Vec<Box<dyn PyrageRecipient>>,
+) -> PyResult<&'p PyBytes> {
+    // This turns each `dyn PyrageRecipient` into a `dyn Recipient`, which
+    // is what the underlying `age` API expects.
+    let recipients = recipients.into_iter().map(|pr| pr.as_recipient()).collect();
+
+    // TODO: More specific exceptions here, rather than ValueError for everything.
+    let encryptor = Encryptor::with_recipients(recipients);
+    let mut encrypted = vec![];
+    let mut writer = encryptor
+        .wrap_output(&mut encrypted)
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+    writer
+        .write_all(plaintext)
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+    writer
+        .finish()
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+
+    // TODO: Avoid this copy. Maybe PyBytes::new_with?
+    Ok(PyBytes::new(py, &encrypted))
+}
+
+#[pyfunction]
+fn decrypt<'p>(
+    py: Python<'p>,
+    ciphertext: &[u8],
+    identities: Vec<Box<dyn PyrageIdentity>>,
+) -> PyResult<&'p PyBytes> {
+    let identities = identities.iter().map(|pi| pi.as_ref().as_identity());
+
+    let decryptor =
+        match age::Decryptor::new(ciphertext).map_err(|e| PyValueError::new_err(e.to_string()))? {
+            age::Decryptor::Recipients(d) => d,
+            age::Decryptor::Passphrase(_) => {
+                return Err(PyValueError::new_err(
+                    "invalid ciphertext (encrypted with passphrase, not identities)",
+                ))
+            }
+        };
+
+    let mut decrypted = vec![];
+    let mut reader = decryptor
+        .decrypt(identities)
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+    reader
+        .read_to_end(&mut decrypted)
+        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+
+    // TODO: Avoid this copy. Maybe PyBytes::new_with?
+    Ok(PyBytes::new(py, &decrypted))
+}
+
 #[pymodule]
 fn pyrage(py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_submodule(x25519::x25519(py)?)
+    m.add_submodule(x25519::x25519(py)?)?;
+
+    m.add_wrapped(wrap_pyfunction!(encrypt))?;
+    m.add_wrapped(wrap_pyfunction!(decrypt))?;
+
+    Ok(())
 }
diff --git a/src/x25519.rs b/src/x25519.rs
index 50f6888..5f9a6be 100644
--- a/src/x25519.rs
+++ b/src/x25519.rs
@@ -4,7 +4,8 @@ use age::secrecy::ExposeSecret;
 use pyo3::{exceptions::PyValueError, prelude::*, types::PyType};
 
 #[pyclass]
-struct Recipient(age::x25519::Recipient);
+#[derive(Clone)]
+pub(crate) struct Recipient(pub(crate) age::x25519::Recipient);
 
 #[pymethods]
 impl Recipient {
@@ -15,13 +16,14 @@ impl Recipient {
             .map_err(PyValueError::new_err)
     }
 
-    fn to_string(&self) -> String {
+    fn __str__(&self) -> String {
         self.0.to_string()
     }
 }
 
 #[pyclass]
-struct Identity(age::x25519::Identity);
+#[derive(Clone)]
+pub(crate) struct Identity(pub(crate) age::x25519::Identity);
 
 #[pymethods]
 impl Identity {
@@ -30,13 +32,13 @@ impl Identity {
         Self(age::x25519::Identity::generate())
     }
 
-    fn to_string(&self) -> String {
-        self.0.to_string().expose_secret().into()
-    }
-
     fn to_public(&self) -> Recipient {
         Recipient(self.0.to_public())
     }
+
+    fn __str__(&self) -> String {
+        self.0.to_string().expose_secret().into()
+    }
 }
 
 pub(crate) fn x25519(py: Python) -> PyResult<&PyModule> {
@@ -45,5 +47,5 @@ pub(crate) fn x25519(py: Python) -> PyResult<&PyModule> {
     module.add_class::<Recipient>()?;
     module.add_class::<Identity>()?;
 
-    Ok(&module)
+    Ok(module)
 }


Test patch:
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/test_pyrage.py b/test/test_pyrage.py
new file mode 100644
index 0000000..cd2995f
--- /dev/null
+++ b/test/test_pyrage.py
@@ -0,0 +1,18 @@
+import unittest
+
+import pyrage
+
+
+class TestPyrage(unittest.TestCase):
+    def test_roundtrip(self):
+        identity = pyrage.x25519.Identity.generate()
+        recipient = identity.to_public()
+
+        encrypted = pyrage.encrypt(b'test', [recipient])
+        decrypted = pyrage.decrypt(encrypted, [identity])
+
+        self.assertEqual(b'test', decrypted)
+
+
+if __name__ == '__main__':
+    unittest.main()