server: serve should be a context manager

sshagent/server.py

@@ -68,7 +68,23 @@ def spawn(func, **kwargs):
     t.join()
 
 
-def run(command, environ):
+@contextlib.contextmanager
+def serve(key_files, signer, sock_path=None):
+    if sock_path is None:
+        sock_path = tempfile.mktemp(prefix='ssh-agent-')
+
+    keys = [formats.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(server_thread, server=server, keys=keys, signer=signer):
+            try:
+                yield environ
+            finally:
+                log.debug('closing server')
+                server.shutdown(socket.SHUT_RD)
+
+
+def run_process(command, environ):
     log.debug('running %r with %r', command, environ)
     env = dict(os.environ)
     env.update(environ)
@@ -80,19 +96,3 @@ def run(command, environ):
     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 = [formats.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(server_thread, server=server, keys=keys, signer=signer):
-            try:
-                ret = run(command=command, environ=environ)
-            finally:
-                log.debug('closing server')
-                server.shutdown(socket.SHUT_RD)
-    return ret

sshagent/trezor_agent.py

@@ -22,8 +22,7 @@ def main():
     level = verbosity[min(args.verbose, len(verbosity) - 1)]
     logging.basicConfig(level=level, format=fmt)
 
-    client = trezor.Client(factory=trezor.TrezorLibrary)
+    with trezor.Client(factory=trezor.TrezorLibrary) as client:
-
         key_files = []
         for label in args.labels:
             pubkey = client.get_public_key(label=label)
@@ -36,16 +35,11 @@ def main():
 
         signer = client.sign_ssh_challenge
 
-    ret = -1
         try:
-        ret = server.serve(
+            with server.serve(key_files=key_files, signer=signer) as env:
-            key_files=key_files,
+                return server.run_process(command=args.command, environ=env)
-            command=args.command,
-            signer=signer)
-        log.info('exitcode: %d', ret)
         except KeyboardInterrupt:
             log.info('server stopped')
-    sys.exit(ret)
 
 if __name__ == '__main__':
-    main()
+    sys.exit(main())