Add basic message ops

conftest.py

@@ -35,6 +35,14 @@ def environment(scoped_instance: GPG) -> GPG:
     return scoped_instance
 
 
+@pytest.fixture(scope="module")
+def smallenv(scoped_instance: GPG):
+    key = scoped_instance.keys.generate_key(
+        "user", email="[email protected]", passphrase="user"
+    )
+    return scoped_instance, key
+
+
 @pytest.fixture
 def interactive(instance: GPG):
     signee = instance.keys.generate_key(

docs/api/operators/messages.md

@@ -0,0 +1,5 @@
+# Messages Operator
+
+Performs operations on message data, such as encryption, decryption, signing, etc
+
+::: gpyg.operators.MessageOperator

gpyg/__init__.py

@@ -1,3 +1,3 @@
 from .util import *
-from .gpg import GPG
-from .models import *
+from .gpg import GPG, Key, KeyOperator, KeyEditor, MessageOperator
+from .models import *

gpyg/gpg.py

@@ -39,3 +39,12 @@ def keys(self) -> KeyOperator:
             KeyOperator: The KeyOperator
         """
         return KeyOperator(self)
+
+    @property
+    def messages(self) -> MessageOperator:
+        """Creates a MessageOperator for this instance
+
+        Returns:
+            MessageOperator: The MessageOperator
+        """
+        return MessageOperator(self)

gpyg/operators/__init__.py

@@ -1 +1,2 @@
 from .keys import KeyOperator, Key, KeyEditor
+from .messages import MessageOperator

gpyg/operators/messages.py

@@ -0,0 +1,113 @@
+from tempfile import NamedTemporaryFile
+from typing import Literal
+from .common import BaseOperator
+from .keys import Key
+from ..util import ExecutionError
+
+
+class MessageOperator(BaseOperator):
+    def encrypt(
+        self,
+        data: bytes,
+        *recipients: Key | str,
+        compress: bool = True,
+        format: Literal["ascii", "pgp"] = "ascii",
+    ) -> bytes:
+        """Encrypt a message to at least one recipient
+
+        Args:
+            data (bytes): Data to encrypt
+            compress (bool, optional): Whether to compress data. Defaults to True.
+            format (ascii | pgp, optional): What format to output. Defaults to "ascii".
+
+        Raises:
+            ValueError: If no recipients were specified
+
+        Returns:
+            bytes: Encrypted data
+        """
+        if len(recipients) == 0:
+            raise ValueError("Must specify at least one recipient")
+        parsed_recipients = " ".join(
+            [f"-r {r.key_id if isinstance(r, Key) else r}" for r in recipients]
+        )
+        cmd = (
+            "gpg {compress} --batch --encrypt {recipients} {armored} --output -".format(
+                compress="-z 0" if compress else "",
+                recipients=parsed_recipients,
+                armored="--armor" if format == "ascii" else "",
+            )
+        )
+        result = self.session.run(cmd, decode=False, input=data)
+        if result.code == 0:
+            return result.output
+        raise ExecutionError(f"Failed to encrypt:\n{result.output}")
+
+    def decrypt(self, data: bytes, key: Key, passphrase: str | None = None) -> bytes:
+        """Decrypt PGP-encrypted data
+
+        Args:
+            data (bytes): Data to decrypt
+            key (Key): Recipient key
+            passphrase (str | None, optional): Recipient passphrase, if present. Defaults to None.
+
+        Raises:
+            ExecutionError: If the operation fails
+
+        Returns:
+            bytes: Decrypted data (with header info removed)
+        """
+        with NamedTemporaryFile() as datafile:
+            datafile.write(data)
+            datafile.seek(0)
+            cmd = f"gpg -u {key.fingerprint} --batch --pinentry-mode loopback --passphrase-fd 0 --output - --decrypt {datafile.name}"
+            result = self.session.run(
+                cmd, decode=False, input=passphrase + "\n" if passphrase else None
+            )
+        if result.code == 0:
+            return result.output.split(b"\n", maxsplit=2)[-1]
+        raise ExecutionError(f"Failed to decrypt:\n{result.output}")
+
+    def get_recipients(
+        self,
+        data: bytes,
+        translate: bool = True,
+        include: list[Literal["known", "unknown"]] = ["known", "unknown"],
+    ) -> list[Key | str]:
+        """Gets all recipients associated with an encrypted message
+
+        Args:
+            data (bytes): Encrypted message
+            translate (bool, optional): Whether to find existing keys
+            include (list[known | unknown], optional): Which keys to include (keys that are known vs keys that are not). Defaults to ["known", "unknown"].
+
+        Raises:
+            ExecutionError: If operation fails
+
+        Returns:
+            list[Key | str]: List of Key objects or, if none match, key IDs
+        """
+        with NamedTemporaryFile() as datafile:
+            datafile.write(data)
+            datafile.seek(0)
+            cmd = f"gpg -d --list-only -v {datafile.name}"
+            result = self.session.run(cmd)
+        if result.code == 0:
+            key_ids = [
+                i.split()[-1] for i in result.output.split("\n") if "public key is" in i
+            ]
+            if translate:
+                keys = []
+                for i in key_ids:
+                    existing = self.gpg.keys.get_key(i)
+                    if i:
+                        if "known" in include:
+                            keys.append(existing)
+                    else:
+                        if "unknown" in include:
+                            keys.append(i)
+
+                return keys
+            else:
+                return key_ids
+        raise ExecutionError(f"Failed to get recipients:\n{result.output}")

gpyg/util/process.py

@@ -144,7 +144,7 @@ def run(
         working_directory: str | None = None,
         timeout: int | None = None,
         decode: bool = True,
-        input: str | None = None,
+        input: str | bytes | None = None,
     ) -> Process:
         options = self.make_kwargs(shell=shell, env=environment, cwd=working_directory)
         parsed_command = self.parse_cmd(
@@ -163,7 +163,9 @@ def run(
         )
 
         if input:
-            self.processes[popen.pid].write(input.encode())
+            self.processes[popen.pid].write(
+                input.encode() if type(input) == str else input
+            )
 
         self.processes[popen.pid].wait(timeout=timeout, kill_on_timeout=True)
         return self.processes[popen.pid]

mkdocs.yml

@@ -8,6 +8,7 @@ nav:
     - GPG Instance: "api/gpg-instance.md"
     - Operators:
       - Key Operations: "api/operators/keys.md"
+      - Message Operations: "api/operators/messages.md"
 
 watch:
   - gpyg

scratchpad.py

@@ -11,14 +11,11 @@
 with TemporaryDirectory(dir="tmp", delete=False) as tmpdir:
     gpg = GPG(homedir=tmpdir, kill_existing_agent=True)
     key = gpg.keys.generate_key("Bongus", passphrase="test")
-    with open(os.path.join(tmpdir, "exported.asc"), "wb") as f:
-        f.write(key.export())
-
-    key.delete()
-    print(gpg.keys.list_keys())
-
-    gpg.keys.import_key(os.path.join(tmpdir, "exported.asc"))
-    print(gpg.keys.list_keys())
+    encrypted = gpg.messages.encrypt(b"A secret message", key, format="pgp")
+    print(encrypted)
+    decrypted = gpg.messages.decrypt(encrypted, key, passphrase="test")
+    print(decrypted)
+    print(gpg.messages.get_recipients(encrypted))
 
     """gpg = GPG(homedir=tmpdir, kill_existing_agent=True)
     key = gpg.keys.generate_key("Bongus", passphrase="test")

tests/test_messages.py

@@ -0,0 +1,22 @@
+from gpyg import *
+
+
+def test_encryption(smallenv):
+    env, key = smallenv
+    DATA = b"test-data"
+    encrypted = env.messages.encrypt(b"test-data", key)
+    assert encrypted != DATA
+    decrypted = env.messages.decrypt(encrypted, key, passphrase="user")
+    assert decrypted == DATA
+
+
+def test_recipients(smallenv):
+    env, key = smallenv
+    DATA = b"test-data"
+    encrypted = env.messages.encrypt(b"test-data", key)
+    assert encrypted != DATA
+
+    recipients = env.messages.get_recipients(encrypted)
+    assert len(recipients) == 1
+    assert isinstance(recipients[0], Key)
+    assert recipients[0].key_id == key.key_id