From 02c8e729b73c450d1795ae07a0e90f0c108218d0 Mon Sep 17 00:00:00 2001 From: Roman Zeyde Date: Tue, 25 Apr 2017 20:43:19 +0300 Subject: [PATCH] ssh: retrieve all keys using a single device session --- trezor_agent/__main__.py | 2 +- trezor_agent/client.py | 18 ++++++++++-------- trezor_agent/tests/test_client.py | 2 +- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/trezor_agent/__main__.py b/trezor_agent/__main__.py index 263fd91..4b31e02 100644 --- a/trezor_agent/__main__.py +++ b/trezor_agent/__main__.py @@ -152,7 +152,7 @@ class JustInTimeConnection(object): def _public_keys(self): """Return a list of SSH public keys (in textual format).""" conn = self.conn_factory() - return [conn.get_public_key(i) for i in self.identities] + return conn.export_public_keys(self.identities) def parse_public_keys(self): """Parse SSH public keys into dictionaries.""" diff --git a/trezor_agent/client.py b/trezor_agent/client.py index 30dfb40..98b5c8c 100644 --- a/trezor_agent/client.py +++ b/trezor_agent/client.py @@ -18,15 +18,17 @@ class Client(object): """Connect to hardware device.""" self.device = device - def get_public_key(self, identity): - """Get SSH public key from the device.""" + def export_public_keys(self, identities): + """Export SSH public keys from the device.""" + public_keys = [] with self.device: - pubkey = self.device.pubkey(identity) - - vk = formats.decompress_pubkey(pubkey=pubkey, - curve_name=identity.curve_name) - return formats.export_public_key(vk=vk, - label=str(identity)) + for i in identities: + pubkey = self.device.pubkey(identity=i) + vk = formats.decompress_pubkey(pubkey=pubkey, + curve_name=i.curve_name) + public_keys.append(formats.export_public_key(vk=vk, + label=str(i))) + return public_keys def sign_ssh_challenge(self, blob, identity): """Sign given blob using a private key on the device.""" diff --git a/trezor_agent/tests/test_client.py b/trezor_agent/tests/test_client.py index 3986cfb..ec90b8e 100644 --- a/trezor_agent/tests/test_client.py +++ b/trezor_agent/tests/test_client.py @@ -49,7 +49,7 @@ def test_ssh_agent(): identity = device.interface.Identity(identity_str='localhost:22', curve_name=CURVE) c = client.Client(device=MockDevice()) - assert c.get_public_key(identity) == PUBKEY_TEXT + assert c.export_public_keys([identity]) == [PUBKEY_TEXT] signature = c.sign_ssh_challenge(blob=BLOB, identity=identity) key = formats.import_public_key(PUBKEY_TEXT)