From 4be55156ed0fda766fe27658f883240c35e38649 Mon Sep 17 00:00:00 2001
From: Roman Zeyde <romanzeyde@google.com>
Date: Fri, 27 May 2016 11:19:10 +0300
Subject: [PATCH] gpg: refactor pubkeys' parsing code

---
 trezor_agent/gpg/decode.py  |  3 ++-
 trezor_agent/gpg/encode.py  |  8 ++++----
 trezor_agent/gpg/keyring.py | 15 ++++++---------
 trezor_agent/gpg/signer.py  |  9 ++++-----
 4 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/trezor_agent/gpg/decode.py b/trezor_agent/gpg/decode.py
index a56efbc..1d771c6 100644
--- a/trezor_agent/gpg/decode.py
+++ b/trezor_agent/gpg/decode.py
@@ -271,8 +271,9 @@ def digest_packets(packets):
     return hashlib.sha256(data_to_hash.getvalue()).digest()
 
 
-def load_public_key(stream, use_custom=False):
+def load_public_key(pubkey_bytes, use_custom=False):
     """Parse and validate GPG public key from an input stream."""
+    stream = io.BytesIO(pubkey_bytes)
     packets = list(parse_packets(util.Reader(stream)))
     pubkey, userid, signature = packets[:3]
     packets = packets[3:]
diff --git a/trezor_agent/gpg/encode.py b/trezor_agent/gpg/encode.py
index 55cd01a..fd47199 100644
--- a/trezor_agent/gpg/encode.py
+++ b/trezor_agent/gpg/encode.py
@@ -2,7 +2,7 @@
 import logging
 import time
 
-from . import keyring, proto
+from . import decode, keyring, proto
 from .. import client, factory, formats, util
 
 log = logging.getLogger(__name__)
@@ -136,10 +136,10 @@ class Factory(object):
         sign_packet = proto.packet(tag=2, blob=signature)
         return pubkey_packet + user_id_packet + sign_packet
 
-    def create_subkey(self):
+    def create_subkey(self, primary_bytes):
         """Export new subkey to `self.user_id` GPG primary key."""
         subkey_packet = proto.packet(tag=14, blob=self.pubkey.data())
-        primary = keyring.get_public_key(self.user_id)
+        primary = decode.load_public_key(primary_bytes)
         log.info('adding subkey to primary GPG key "%s" (%s)',
                  self.user_id, util.hexlify(primary['key_id']))
         data_to_sign = primary['_to_hash'] + self.pubkey.data_to_hash()
@@ -176,7 +176,7 @@ class Factory(object):
             hashed_subpackets=hashed_subpackets,
             unhashed_subpackets=unhashed_subpackets)
         sign_packet = proto.packet(tag=2, blob=signature)
-        return subkey_packet + sign_packet
+        return primary_bytes + subkey_packet + sign_packet
 
     def sign_message(self, msg, sign_time=None):
         """Sign GPG message at specified time."""
diff --git a/trezor_agent/gpg/keyring.py b/trezor_agent/gpg/keyring.py
index 427c42c..bdf5c67 100644
--- a/trezor_agent/gpg/keyring.py
+++ b/trezor_agent/gpg/keyring.py
@@ -8,7 +8,6 @@ import re
 import socket
 import subprocess
 
-from . import decode
 from .. import util
 
 log = logging.getLogger(__name__)
@@ -45,7 +44,7 @@ def _recvline(sock):
 
 
 def unescape(s):
-    """Unescape ASSUAN message."""
+    """Unescape ASSUAN message (0xAB <-> '%AB')."""
     s = bytearray(s)
     i = 0
     while i < len(s):
@@ -151,13 +150,11 @@ def get_keygrip(user_id):
     return re.findall(r'Keygrip = (\w+)', output)[0]
 
 
-def get_public_key(user_id, use_custom=False):
-    """Load existing GPG public key for `user_id` from local keyring."""
+def export_public_key(user_id):
+    """Export GPG public key for specified `user_id`."""
     args = ['gpg2', '--export'] + ([user_id] if user_id else [])
-    pubkey_bytes = subprocess.check_output(args=args)
-    if pubkey_bytes:
-        return decode.load_public_key(io.BytesIO(pubkey_bytes),
-                                      use_custom=use_custom)
-    else:
+    result = subprocess.check_output(args=args)
+    if not result:
         log.error('could not find public key %r in local GPG keyring', user_id)
         raise KeyError(user_id)
+    return result
diff --git a/trezor_agent/gpg/signer.py b/trezor_agent/gpg/signer.py
index 214beb4..c5cb29e 100755
--- a/trezor_agent/gpg/signer.py
+++ b/trezor_agent/gpg/signer.py
@@ -3,7 +3,6 @@
 import argparse
 import contextlib
 import logging
-import subprocess as sp
 import sys
 import time
 import os
@@ -21,9 +20,8 @@ def run_create(args):
 
     with contextlib.closing(f):
         if args.subkey:
-            subkey = f.create_subkey()
-            primary = sp.check_output(['gpg2', '--export', user_id])
-            result = primary + subkey
+            primary_key = keyring.export_public_key(user_id=user_id)
+            result = f.create_subkey(primary_bytes=primary_key)
         else:
             result = f.create_primary()
 
@@ -32,7 +30,8 @@ def run_create(args):
 
 def run_sign(args):
     """Generate a GPG signature using hardware-based device."""
-    pubkey = keyring.get_public_key(user_id=None, use_custom=True)
+    pubkey = decode.load_public_key(keyring.export_public_key(user_id=None),
+                                    use_custom=True)
     f = encode.Factory.from_public_key(pubkey=pubkey,
                                        user_id=pubkey['user_id'])
     with contextlib.closing(f):