Support "fast-path" key listing

https://dev.gnupg.org/rG40da61b89b62dcb77847dc79eb159e885f52f817#change-o4DEJvEV1Dx2

Also, refactor decoding and add a few tests.

libagent/gpg/agent.py

@@ -99,7 +99,7 @@ class Handler:
             b'SETHASH': lambda _, args: self.set_hash(*args),
             b'PKSIGN': lambda conn, _: self.pksign(conn),
             b'PKDECRYPT': lambda conn, _: self.pkdecrypt(conn),
-            b'HAVEKEY': lambda _, args: self.have_key(*args),
+            b'HAVEKEY': lambda conn, args: self.have_key(conn, *args),
             b'KEYINFO': _key_info,
             b'SCD': self.handle_scd,
             b'GET_PASSPHRASE': self.handle_get_passphrase,
@@ -198,8 +198,16 @@ class Handler:
         ec_point = self.client.ecdh(identity=identity, pubkey=remote_pubkey)
         keyring.sendline(conn, b'D ' + _serialize_point(ec_point))
 
-    def have_key(self, *keygrips):
+    def have_key(self, conn, *keygrips):
         """Check if any keygrip corresponds to a TREZOR-based key."""
+        if len(keygrips) == 1 and keygrips[0].startswith(b"--list="):
+            # Support "fast-path" key listing:
+            # https://dev.gnupg.org/rG40da61b89b62dcb77847dc79eb159e885f52f817
+            keygrips = list(decode.iter_keygrips(pubkey_bytes=self.pubkey_bytes))
+            log.debug('keygrips: %r', keygrips)
+            keyring.sendline(conn, b'D ' + util.assuan_serialize(b''.join(keygrips)))
+            return
+
         for keygrip in keygrips:
             try:
                 self.get_identity(keygrip=keygrip)

libagent/gpg/decode.py

@@ -282,18 +282,20 @@ HASH_ALGORITHMS = {
 }
 
 
-def load_by_keygrip(pubkey_bytes, keygrip):
-    """Return public key and first user ID for specified keygrip."""
+def _parse_pubkey_packets(pubkey_bytes):
     stream = io.BytesIO(pubkey_bytes)
-    packets = list(parse_packets(stream))
     packets_per_pubkey = []
-    for p in packets:
+    for p in parse_packets(stream):
         if p['type'] == 'pubkey':
             # Add a new packet list for each pubkey.
             packets_per_pubkey.append([])
         packets_per_pubkey[-1].append(p)
+    return packets_per_pubkey
 
-    for packets in packets_per_pubkey:
+
+def load_by_keygrip(pubkey_bytes, keygrip):
+    """Return public key and first user ID for specified keygrip."""
+    for packets in _parse_pubkey_packets(pubkey_bytes):
         user_ids = [p for p in packets if p['type'] == 'user_id']
         for p in packets:
             if p.get('keygrip') == keygrip:
@@ -301,6 +303,15 @@ def load_by_keygrip(pubkey_bytes, keygrip):
     raise KeyError('{} keygrip not found'.format(util.hexlify(keygrip)))
 
 
+def iter_keygrips(pubkey_bytes):
+    """Iterate over all keygrips in this pubkey."""
+    for packets in _parse_pubkey_packets(pubkey_bytes):
+        for p in packets:
+            keygrip = p.get('keygrip')
+            if keygrip:
+                yield keygrip
+
+
 def load_signature(stream, original_data):
     """Load signature from stream, and compute GPG digest for verification."""
     signature, = list(parse_packets((stream)))

libagent/gpg/tests/test_decode.py

@@ -1,6 +1,5 @@
-import glob
 import io
-import os
+import pathlib
 
 import pytest
 
@@ -30,8 +29,8 @@ def test_mpi():
     assert decode.parse_mpis(util.Reader(s), n=2) == [0x123, 5]
 
 
-cwd = os.path.join(os.path.dirname(__file__))
-input_files = glob.glob(os.path.join(cwd, '*.gpg'))
+cwd = pathlib.Path(__file__).parent
+input_files = cwd.glob('*.gpg')
 
 
 @pytest.fixture(params=input_files)
@@ -60,3 +59,20 @@ def test_has_custom_subpacket():
 def test_load_by_keygrip_missing():
     with pytest.raises(KeyError):
         decode.load_by_keygrip(pubkey_bytes=b'', keygrip=b'')
+
+
+def test_keygrips():
+    pubkey_bytes = (cwd / "romanz-pubkey.gpg").open("rb").read()
+    keygrips = list(decode.iter_keygrips(pubkey_bytes))
+    assert [k.hex() for k in keygrips] == [
+        '7b2497258d76bc6539ed88d018cd1c739e2dbb6c',
+        '30ae97f3d8e0e34c5ed80e1715fd442ca24c0a8e',
+    ]
+
+    for keygrip in keygrips:
+        pubkey_dict, user_ids = decode.load_by_keygrip(pubkey_bytes, keygrip)
+        assert pubkey_dict['keygrip'] == keygrip
+        assert [u['value'] for u in user_ids] == [
+            b'Roman Zeyde <roman.zeyde@gmail.com>',
+            b'Roman Zeyde <me@romanzey.de>',
+        ]