From bd22d7512b4a87b7ae097a623659c0cea1ca9f2c Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Sat, 4 Jul 2015 21:15:09 +0300
Subject: [PATCH] trezor: refactor label parsing and handling

---
 sshagent/trezor.py       | 52 +++++++++++++++++++---------------------
 sshagent/trezor_agent.py |  4 +++-
 2 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/sshagent/trezor.py b/sshagent/trezor.py
index 228433e..a15c686 100644
--- a/sshagent/trezor.py
+++ b/sshagent/trezor.py
@@ -1,4 +1,5 @@
 import io
+import re
 import struct
 import binascii
 
@@ -63,8 +64,8 @@ class Client(object):
         pubkey = node.node.public_key
         return formats.export_public_key(pubkey=pubkey, label=label)
 
-    def sign_ssh_challenge(self, label, blob):
-        identity = self.factory.parse_identity(label)
+    def sign_ssh_challenge(self, identity, blob):
+        label = _identity_to_string(identity)
         msg = _parse_ssh_blob(blob)
 
         log.info('confirm user %s connection to %r using Trezor...',
@@ -82,40 +83,34 @@ class Client(object):
         return (r, s)
 
 
-def _lsplit(s, sep):
-    p = None
-    if sep in s:
-        p, s = s.split(sep, 1)
-    return (p, s)
-
+def _split(s, *parts):
+    return re.match(regexp, s).groups()
 
-def _rsplit(s, sep):
-    p = None
-    if sep in s:
-        s, p = s.rsplit(sep, 1)
-    return (s, p)
 
+_parts = [
+    '^'
+    r'(?:(?P<proto>.*)://)?',
+    r'(?:(?P<user>.*)@)?',
+    r'(?P<host>.*?)',
+    r'(?::(?P<port>\w*))?',
+    r'(?:/(?P<path>.*))?',
+    '$'
+]
+_identity_regexp = re.compile(''.join(_parts))
 
 def _string_to_identity(s):
-    proto, s = _lsplit(s, '://')
-    user, s = _lsplit(s, '@')
-    s, path = _rsplit(s, '/')
-    host, port = _rsplit(s, ':')
-
-    if not proto:
-        proto = 'ssh'  # otherwise, Trezor will use SECP256K1 curve
+    m = _identity_regexp.match(s)
+    result = m.groupdict()
+    if not result.get('proto'):
+        result['proto'] = 'ssh'  # otherwise, Trezor will use SECP256K1 curve
 
-    result = [
-        ('proto', proto), ('user', user), ('host', host),
-        ('port', port), ('path', path)
-    ]
-    return {k: v for k, v in result if v}
+    log.debug('parsed identity: %s', result)
+    return {k: v for k, v in result.items() if v}
 
 
 def _identity_to_string(identity):
-    result = []
-    if identity.proto:
-        result.append(identity.proto + '://')
+    assert identity.proto == 'ssh'
+    result = [identity.proto + '://']
     if identity.user:
         result.append(identity.user + '@')
     result.append(identity.host)
@@ -149,6 +144,7 @@ def _parse_ssh_blob(data):
         i.read(1)  # TBD
         res['key_type'] = util.read_frame(i)
         res['pubkey'] = util.read_frame(i)
+        assert not i.read()
         log.debug('%s: user %r via %r (%r)',
                   res['conn'], res['user'], res['auth'], res['key_type'])
         log.debug('nonce: %s', binascii.hexlify(res['nonce']))
diff --git a/sshagent/trezor_agent.py b/sshagent/trezor_agent.py
index cea9eae..0f5c709 100644
--- a/sshagent/trezor_agent.py
+++ b/sshagent/trezor_agent.py
@@ -53,7 +53,9 @@ def main():
             sys.stdout.write(public_key)
             return
 
-        signer = client.sign_ssh_challenge
+        def signer(label, blob):
+            identity = client.get_identity(label=label)
+            return client.sign_ssh_challenge(identity=identity, blob=blob)
 
         try:
             with server.serve(public_keys=[public_key], signer=signer) as env: