add support Trezor SSH agent

agent.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python
+import socket
+import os
+import sys
+import subprocess
+import argparse
+import tempfile
+import contextlib
+import threading
+import logging
+log = logging.getLogger(__name__)
+
+import protocol
+import trezor
+
+
+def load_keys(key_files):
+    keys = []
+    for f in key_files:
+        k = protocol.load_public_key(f)
+        keys.append(k)
+    return keys
+
+
+@contextlib.contextmanager
+def unix_domain_socket_server(sock_path):
+    log.debug('serving on SSH_AUTH_SOCK=%s', sock_path)
+    try:
+        os.remove(sock_path)
+    except OSError:
+        if os.path.exists(sock_path):
+            raise
+
+    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    server.bind(sock_path)
+    server.listen(1)
+    try:
+        yield server
+    finally:
+        os.remove(sock_path)
+
+
+def worker_thread(server, keys, signer):
+    log.debug('worker thread started')
+    while True:
+        log.debug('waiting for connection on %s', server.getsockname())
+        try:
+            conn, _ = server.accept()
+        except socket.error as e:
+            log.debug('server error: %s', e, exc_info=True)
+            break
+        with contextlib.closing(conn):
+            protocol.handle_connection(conn, keys, signer)
+    log.debug('worker thread stopped')
+
+
+@contextlib.contextmanager
+def spawn(func, **kwargs):
+    t = threading.Thread(target=func, kwargs=kwargs)
+    t.start()
+    yield
+    t.join()
+
+
+def run(command, environ):
+    log.debug('running %r with %r', command, environ)
+    env = dict(os.environ)
+    env.update(environ)
+    p = subprocess.Popen(args=command, env=env)
+    log.debug('subprocess %d is running', p.pid)
+    ret = p.wait()
+    log.debug('subprocess %d exited: %d', p.pid, ret)
+    return ret
+
+
+def serve(key_files, command, signer, sock_path=None):
+    if sock_path is None:
+        sock_path = tempfile.mktemp(prefix='ssh-agent-')
+
+    keys = [protocol.parse_public_key(k) for k in key_files]
+    environ = {'SSH_AUTH_SOCK': sock_path, 'SSH_AGENT_PID': str(os.getpid())}
+    with unix_domain_socket_server(sock_path) as server:
+        with spawn(worker_thread, server=server, keys=keys, signer=signer):
+            ret = run(command=command, environ=environ)
+            log.debug('closing server')
+            server.shutdown(socket.SHUT_RD)
+
+    log.info('exitcode: %d', ret)
+    sys.exit(ret)
+
+
+def main():
+    fmt = '%(asctime)s %(levelname)-12s %(message)-100s [%(filename)s]'
+    p = argparse.ArgumentParser()
+    p.add_argument('-k', '--key-label', dest='labels', action='append')
+    p.add_argument('-v', '--verbose', action='count', default=0)
+    p.add_argument('command', type=str, nargs='*')
+    args = p.parse_args()
+
+    verbosity = [logging.WARNING, logging.INFO, logging.DEBUG]
+    level = verbosity[min(args.verbose, len(verbosity))]
+    logging.basicConfig(level=level, format=fmt)
+
+    client = trezor.Client()
+
+    key_files = []
+    for label in args.labels:
+        pubkey = client.get_public_key(label=label)
+        key_files.append(trezor.export_public_key(pubkey=pubkey, label=label))
+
+    if not args.command:
+        sys.stdout.write(''.join(key_files))
+        return
+
+    signer = client.sign_ssh_challenge
+
+    try:
+        serve(key_files=key_files, command=args.command, signer=signer)
+    except KeyboardInterrupt:
+        log.info('server stopped')
+
+if __name__ == '__main__':
+    main()