From b9b7b8dafd8929d45fb1d6672045e97c12066d0d Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Fri, 22 Jul 2016 22:51:02 +0300
Subject: [PATCH] gpg: re-structure public key packets for easier parsing

---
 trezor_agent/gpg/decode.py            | 25 +++++++++++++++++++++++++
 trezor_agent/gpg/tests/test_decode.py |  3 ++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/trezor_agent/gpg/decode.py b/trezor_agent/gpg/decode.py
index 7f7be3c..6035142 100644
--- a/trezor_agent/gpg/decode.py
+++ b/trezor_agent/gpg/decode.py
@@ -1,5 +1,6 @@
 """Decoders for GPG v2 data structures."""
 import base64
+import copy
 import functools
 import hashlib
 import io
@@ -272,6 +273,30 @@ def digest_packets(packets, hasher):
     return hasher.digest()
 
 
+def collect_packets(packets, types_to_collect):
+    """Collect specified packet types into their leading packet."""
+    packet = None
+    result = []
+    for p in packets:
+        if p['type'] in types_to_collect:
+            packet.setdefault(p['type'], []).append(p)
+        else:
+            packet = copy.deepcopy(p)
+            result.append(packet)
+    return result
+
+
+def parse_public_key(stream):
+    """Parse GPG public key into hierarchy of packets."""
+    packets = list(parse_packets(stream))
+    packets = collect_packets(packets, {'signature'})
+    packets = collect_packets(packets, {'user_id', 'user_attribute'})
+    packets = collect_packets(packets, {'subkey'})
+    if len(packets) != 1:
+        raise ValueError('Unexpected packet structure: %r', packets)
+    return packets[0]
+
+
 HASH_ALGORITHMS = {
     1: 'md5',
     2: 'sha1',
diff --git a/trezor_agent/gpg/tests/test_decode.py b/trezor_agent/gpg/tests/test_decode.py
index 27f7383..92fc82d 100644
--- a/trezor_agent/gpg/tests/test_decode.py
+++ b/trezor_agent/gpg/tests/test_decode.py
@@ -110,4 +110,5 @@ def public_key_path(request):
 
 
 def test_gpg_files(public_key_path):  # pylint: disable=redefined-outer-name
-    decode.load_public_key(open(public_key_path, 'rb').read())
+    with open(public_key_path, 'rb') as f:
+        decode.parse_public_key(f)