read_agentproto_idcomments() to read loaded IDs from an SSH agent

This will be used to query an agent for the loaded keys so that we
don't try to reload ones that are already there.

Author: 0cjs

bin/ckssh.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 from argparse       import ArgumentParser
+from binascii       import hexlify
 from collections    import namedtuple as ntup
 from os             import path
 from pathlib        import Path
@@ -13,6 +14,55 @@
 CONFIG_FILE     = '~/.ssh/ckssh_config'
 DEVNULL         = open(os.devnull, 'w')
 
+############################################################
+#   SSH agent protocol functions
+
+class SSHAgentProtoError(RuntimeError):
+    pass
+
+def read_agentproto_int(stream, length):
+    bs = stream.read(length)
+    if len(bs) != length:
+        raise SSHAgentProtoError('Short int: {}'.format(bs))
+    return int.from_bytes(bs, byteorder='big')
+
+def read_agentproto_bstr(stream):
+    length = read_agentproto_int(stream, 4)
+    bs = stream.read(length)
+    if len(bs) != length:
+        raise SSHAgentProtoError('Short string: {}'.format(bs))
+    return bs
+
+def read_agentproto_idcomments(stream):
+    ''' From the given I/O stream, Read and parse an
+        ``SSH2_AGENT_IDENTITIES_ANSWER`` response to an
+        ``SSH2_AGENTC_REQUEST_IDENTITIES`` request.
+        Return a list with the comment for each identity.
+
+        For protocol details see section 2.5.2 of
+        <http://api.libssh.org/rfc/PROTOCOL.agent>.
+    '''
+    #   We don't actually use the message length, instead relying on
+    #   the count of keys and string lengths, but we parse it to make
+    #   sure this isn't a bad message.
+    msglen = read_agentproto_int(stream, 4)
+
+    msgtype = read_agentproto_int(stream, 1)
+    if msgtype != 0xC:
+        raise SSHAgentProtoError(
+            'Unknown message type: {}'.format(msgtype))
+
+    keycount = read_agentproto_int(stream, 4)
+    comments = []
+    for i in range(0, keycount):
+        read_agentproto_bstr(stream)             # key blob
+        bs = read_agentproto_bstr(stream)        # key comment
+        comments.append(bs.decode('ascii'))
+    return comments
+
+    raise SSHAgentProtoError()
+    return []
+
 ############################################################
 #   Compartment classes and functions
 

t/unit.py

@@ -1,12 +1,59 @@
-from    io import StringIO
+from    io import BytesIO, StringIO
+from    binascii import unhexlify
 from    pathlib import Path
 import  pytest
 
 import  ckssh
 from    ckssh import (
+        SSHAgentProtoError, read_agentproto_idcomments,
         runtimedir, Compartment, parseconfig, canexec, evalwrite, CK,
         )
 
+####################################################################
+#   SSH agent protocol
+
+def test_ssh2_ids_answer_decode_0():
+    answer0 = (
+        0, 0, 0, 5,     # message length
+        0xC,            # SSH2_AGENT_IDENTITIES_ANSWER
+        0, 0, 0, 0,     # identities count: 0
+    )
+    assert [] == read_agentproto_idcomments(BytesIO(bytes(answer0)))
+
+def test_ssh2_ids_answer_decode_2():
+    answer2 = (
+        b''
+        b'\0\0\0\5'     # message length
+        b'\x0C'         # SSH2_AGENT_IDENTITIES_ANSWER
+        b'\0\0\0\2'     # identities count: 2
+
+        b'\0\0\0\1'     # 1st answer blob length: 1
+        b'\x7E'
+        b'\0\0\0\3'     # 1st answer comment length: 3
+        b'foo'
+
+        b'\0\0\0\0'     # 2nd answer blob length: 0
+        b''
+        b'\0\0\0\7'     # 2nd answer comment length: 7
+        b'bar baz'
+    )
+    assert ['foo', 'bar baz'] \
+        == read_agentproto_idcomments(BytesIO(bytes(answer2)))
+
+@pytest.mark.parametrize('badmsg', (
+    '',
+    '00000001',     # no response type
+    '0000000102',   # bad response type
+    '000000020C',   # bad length
+    ))
+def test_ssh2_ids_answer_decode_err(badmsg):
+    bs = unhexlify(badmsg)  # Convert printable test param to bytes
+    with pytest.raises(SSHAgentProtoError):
+        read_agentproto_idcomments(BytesIO(bs))
+
+####################################################################
+#   ckssh stuff
+
 TESTDIR = Path(__file__).parent
 CONFIGFILE = Path(TESTDIR, 'mock_home/.ssh/ckssh_config')